import asyncio
from typing import List, Dict, Any
from openai import AsyncOpenAI
from openai import OpenAI


class TextGenerator:
                  

    def __init__(self, config):
        self.config = config
        self.client = OpenAI(api_key=config.api_key, base_url=config.base_url)
        self.async_client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)

        if config.verbose:
            print("TextGenerator initialized")

    def generate(self, query: str, docs: List[Dict[str, Any]]) -> str:
                  
                 
        docs = docs[:3]

                
        docs_text = ""
        for i, doc in enumerate(docs, 1):
            metadata = doc.get('metadata', {})
            title = metadata.get('title', 'No title')

                          
            docs_text += f"=== Reference Material {i} ===\n"
            docs_text += f"Title: {title}\n"
            docs_text += f"Content: {doc['content']}\n\n"

        user_prompt = f"""Context information is below:
            {docs_text}
        Given the context information and not prior knowledge, answer the query.
        Query: {query}
        Answer:
        """

        try:
            response = self.client.chat.completions.create(
                model=self.config.llm_model,
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=self.config.temperature
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            return f"Generation failed: {e}"

    async def generate_async(self, query: str, docs: List[Dict[str, Any]]) -> str:
                    
                 
        docs = docs[:3]

                
        docs_text = ""
        for i, doc in enumerate(docs, 1):
            metadata = doc.get('metadata', {})
            title = metadata.get('title', 'No title')

                          
            docs_text += f"=== Reference Material {i} ===\n"
            docs_text += f"Title: {title}\n"
            docs_text += f"Content: {doc['content']}\n\n"

        user_prompt = f"""Context information is below:
            {docs_text}
        Given the context information and not prior knowledge, answer the query.
        Query: {query}
        Answer:
        """

                
        max_retries = 3
        for attempt in range(max_retries):
            try:
                response = await self.async_client.chat.completions.create(
                    model=self.config.llm_model,
                    messages=[
                        {"role": "system", "content": "You are a helpful assistant."},
                        {"role": "user", "content": user_prompt}
                    ],
                    temperature=self.config.temperature
                )
                return response.choices[0].message.content.strip()
            except Exception as e:
                error_msg = str(e)

                                    
                if attempt == max_retries - 1:
                    if "DOCTYPE html" in error_msg or "<!DOCTYPE" in error_msg:
                        return "Generation failed: API returned HTML error page, possibly service issue or request redirected"
                    elif "400" in error_msg:
                        return f"Generation failed: Request parameter error (400) - {error_msg}"
                    elif "401" in error_msg:
                        return f"Generation failed: Authentication failed (401) - {error_msg}"
                    elif "429" in error_msg:
                        return f"Generation failed: Request rate limit (429) - {error_msg}"
                    elif "500" in error_msg:
                        return f"Generation failed: Server internal error (500) - {error_msg}"
                    else:
                        return f"Generation failed: {error_msg}"
                else:
                                    
                    if self.config.verbose:
                        print(f"Error generating for attempt {attempt}: {error_msg}")
                    await asyncio.sleep(1 * (attempt + 1))          
