"""
LLM Client - Connect to LLM for RAG queries
"""

import os
from typing import Optional, Dict, Union

try:
    from langchain_openai import ChatOpenAI
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_core.output_parsers import StrOutputParser
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False


class LLMClient:
    """LLM client for RAG queries"""
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        model: Optional[str] = None,
        temperature: float = 0.7
    ):
        """
        Initialize LLM client
        
        Args:
            api_key: LLM API key (default: from OPENAI_API_KEY env var)
            base_url: LLM API base URL (default: from OPENAI_BASE_URL env var)
            model: Model name (default: from OPENAI_MODEL env var or 'gpt-3.5-turbo')
            temperature: Temperature parameter (default: 0.7)
        """
        if not LANGCHAIN_AVAILABLE:
            raise ImportError(
                "LangChain is required for LLM functionality. "
                "Install it with: pip install mineru-rag[rag]"
            )
        
        self.api_key = api_key or os.environ.get("OPENAI_API_KEY")
        self.base_url = base_url or os.environ.get("OPENAI_BASE_URL")
        self.model = model or os.environ.get("OPENAI_MODEL", "gpt-3.5-turbo")
        self.temperature = temperature
        
        if not self.api_key:
            raise ValueError(
                "OPENAI_API_KEY is required. Set it as environment variable or pass as parameter."
            )
        
        if not self.base_url:
            raise ValueError(
                "OPENAI_BASE_URL is required. Set it as environment variable or pass as parameter."
            )
        
        # Initialize LLM
        self.llm = ChatOpenAI(
            model=self.model,
            base_url=self.base_url,
            api_key=self.api_key,
            temperature=self.temperature
        )
        
        # Create prompt template
        self.prompt_template = ChatPromptTemplate.from_messages([
            ("system", "You are a helpful assistant that answers questions based on the provided context."),
            ("user", """Based on the following context, please answer the question. 
If the context doesn't contain relevant information, please say so.

Context:
{context}

Question: {question}

Please provide a detailed and accurate answer:""")
        ])
        
        # Create chain
        self.chain = self.prompt_template | self.llm | StrOutputParser()
    
    def query(
        self,
        question: str,
        context: str
    ) -> str:
        """
        Query LLM with context
        
        Args:
            question: Question to ask
            context: Context from RAG retrieval
            
        Returns:
            LLM response
        """
        return self.chain.invoke({
            "question": question,
            "context": context
        })
    
    def query_with_rag(
        self,
        rag_result: Dict
    ) -> Dict:
        """
        Query LLM with RAG result
        
        Args:
            rag_result: Result from RAGBuilder.query()
            
        Returns:
            Dictionary with answer and sources
        """
        question = rag_result['question']
        context = rag_result['context']
        
        answer = self.query(question, context)
        
        return {
            'question': question,
            'answer': answer,
            'sources': rag_result['sources'],
            'num_sources': rag_result['num_sources']
        }

