import logging
from typing import Dict, List, Optional
from langchain_groq import ChatGroq
from langchain.callbacks import StreamingStdOutCallbackHandler
from langchain.callbacks.manager import CallbackManager

# Configure Logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class ThemeSynthesizer:
    """A service for extracting themes and insights from document collections using Groq Llama models."""

    # Class constants for different models available on Groq
    DEFAULT_MODELS = {
        "small": "llama3-8b-8192",
        "medium": "llama3-70b-8192",
        "large": "llama3-70b-8192",
        "mixtral": "mixtral-8x7b-32768",
        "gemma": "gemma-7b-it"
    }

    # Template for individual document analysis
    DOCUMENT_ANALYSIS_TEMPLATE = """
    You are an expert document analyst focused on precise extraction.
    Given the following document content, extract relevant information in response to this query:

    Query: {query}

    DOCUMENT CONTENT:
    {content}

    Your task:
    1. Extract the most relevant information that answers the query
    2. Be concise yet complete in your extraction
    3. Include EXACT citations by page number and paragraph (e.g., "Page 3, Page 2")
    4. Only provide information that is directly mentioned in the document
    5. If no relevant information exists, state "No relevant information found"

    Your answer format must be:
    EXTRACTED ANSWER: [Your concise answer here]
    CITATION: [Exact page and paragraph reference]
    """

    # Template for theme synthesis across documents
    THEME_SYNTHESIS_TEMPLATE = """
    You are the expert research analyst synthesizing findings across multiple documents.
    Below are answers extracted from different documents in response to a specific query.

    Query: {query}

    DOCUMENT ANSWERS:
    {document_answers}

    Your task:
    1. Identify 2-6 key themes or patterns across these documents
    2. For each theme, provide a clear title and brief explanation
    3. Be concise yet complete in your extraction
    4. Mention specific document IDs (e.g, DOC001) that support each theme
    5. Format your response as a chat message that synthesizes the information

    Format your response as follows:

    Theme 1 – [Theme Title]:
    Documents ([Document IDs]) [brief explanation with citations]

    Theme 2 – [Theme Title]:
    Documents ([Document IDs]) [brief explanation with citations]

    ...and so on for each theme.
    """

    def __init__(
        self,
        model_size: str = "small",
        custom_model_name: Optional[str] = None,
        temperature: float = 0.2,
        max_tokens: int = 8192,
        streaming: bool = False,
        api_key: Optional[str] = None
    ):
        """Initialize the document analyzer with configurable parameters."""
        # Check for API key
        if not api_key:
            raise ValueError("Groq API key is required. Get one from https://console.groq.com/")

        # Set up callbacks for streaming if enabled
        callbacks = None
        if streaming:
            callbacks = CallbackManager([StreamingStdOutCallbackHandler()])

        # Select model based on size or custom name
        model_name = custom_model_name or self.DEFAULT_MODELS.get(model_size, self.DEFAULT_MODELS["small"])

        # Initialize the Groq language model
        try:
            self.llm = ChatGroq(
                groq_api_key=api_key,
                model_name=model_name,
                temperature=temperature,
                max_tokens=max_tokens,
                streaming=streaming,
                callbacks=callbacks
            )
            logger.info(f"Initialized Groq model: {model_name}")
        except Exception as e:
            logger.error(f"Failed to initialize Groq model: {e}")
            raise ValueError(f"Failed to initialize Groq model: {e}")

        # Store configuration
        self.config = {
            "model": model_name,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "provider": "groq"
        }

    def analyze_single_document(
        self,
        document: Dict,
        query: str
    ) -> Dict:
        """
        Process a single document to extract relevant information for the query.

        Args:
            document: Document dictionary with 'content' and 'metadata'
            query: User query to answer

        Returns:
            Dictionary with extracted answer and citation
        """
        try:
            doc_id = document.get('metadata', {}).get('doc_id', 'Unknown')
            content = document.get('content', '')

            if not content.strip():
                return {
                    "doc_id": doc_id,
                    "extracted_answer": "No content available in document",
                    "citation": "N/A",
                    "status": "error"
                }

            logger.info(f"Analyzing document {doc_id} for query: {query[:50]}...")

            # Format the prompt with query and content
            formatted_prompt = self.DOCUMENT_ANALYSIS_TEMPLATE.format(
                query=query,
                content=content
            )

            # Run document analysis directly with the LLM
            response = self.llm.invoke(formatted_prompt)

            # Extract the result from the response
            if hasattr(response, 'content'):
                response = response.content
            elif isinstance(response, dict) and "text" in response:
                response = response["text"]

            # Parse response to extract answer and citation
            extracted_answer = "No relevant information found"
            citation = "N/A"

            # Simple parsing - can be made more robust
            if "EXTRACTED ANSWER:" in response:
                answer_parts = response.split("EXTRACTED ANSWER:")
                if len(answer_parts) > 1:
                    citation_parts = answer_parts[1].split("CITATION:")
                    extracted_answer = citation_parts[0].strip()
                    if len(citation_parts) > 1:
                        citation = citation_parts[1].strip()

            return {
                "doc_id": doc_id,
                "extracted_answer": extracted_answer,
                "citation": citation,
                "status": "success"
            }

        except Exception as e:
            logger.error(f"Error analyzing document: {e}", exc_info=True)
            return {
                "doc_id": document.get('metadata', {}).get('doc_id', 'Unknown'),
                "extracted_answer": f"Error during analysis: {str(e)}",
                "citation": "N/A",
                "status": "error"
            }

    def batch_analyze_documents(
            self,
            documents: List[Dict],
            query: str
    ) -> List[Dict]:
        """Process a batch of documents to extract relevant information for the query.

        Args:
            documents: List of document dictionaries with 'content' and 'metadata'
            query: User query to answer

        Returns:
            List of dictionaries with extracted answers and citations.

        """
        try:
            results = []
            for doc in documents:
                # Ensure Document has proper Id
                if 'metadata' not in doc:
                    doc['metadata'] = {}
                if 'doc_id' not in doc['metadata']:
                    doc['metadata']['doc_id'] = f"DOC{len(results)+1:03d}"

                # Analyze document
                result = self.analyze_single_document(doc, query)
                results.append(result)
            return results
        except Exception as e:
            logger.error(f"Error batch analyzing documents: {e}", exc_info=True)
            return []

    def synthesize_themes(
        self,
        document_results: List[Dict],
        query: str
    ) -> Dict:
        """
        Synthesize themes from individual document analysis results.

        Args:
            document_results: Results from batch_analyze_documents
            query: Original user query

        Returns:
            Dictionary with synthesized themes and status
        """
        try:
            if not document_results:
                return {
                    "themes": "No documents provided for analysis.",
                    "status": "warning"
                }

            # Format document answers for the synthesis prompt
            formatted_answers = []
            for res in document_results:
                doc_id = res.get("doc_id", "Unknown")
                answer = res.get("extracted_answer", "No answer available")
                citation = res.get("citation", "N/A")

                formatted_answer = f"Document ID: {doc_id}\nExtracted Answer: {answer}\nCitation: {citation}"
                formatted_answers.append(formatted_answer)

            combined_answers = "\n\n".join(formatted_answers)

            # Format the prompt with query and document answers
            formatted_prompt = self.THEME_SYNTHESIS_TEMPLATE.format(
                query=query,
                document_answers=combined_answers
            )

            # Run theme synthesis directly with the LLM
            response = self.llm.invoke(formatted_prompt)

            # Extract the result from the response
            if hasattr(response, 'content'):
                response = response.content
            elif isinstance(response, dict) and "text" in response:
                response = response["text"]

            return {
                "themes": response.strip(),
                "status": "success",
                "stats": {
                    "document_count": len(document_results),
                    "query": query,
                    "model": self.config["model"]
                }
            }

        except Exception as e:
            logger.error(f"Error synthesizing themes: {e}", exc_info=True)
            return {
                "themes": f"Error during theme synthesis: {str(e)}",
                "status": "error",
                "error": str(e)
            }

    def process_query(
        self,
        documents: List[Dict],
        query: str
    ) -> Dict:
        """
        Complete document processing pipeline - analyze documents and synthesize themes.

        Args:
            documents: List of document dictionaries
            query: User query to answer

        Returns:
            Dictionary with individual document results and synthesized themes
        """
        # Step 1: Analyze each document individually
        document_results = self.batch_analyze_documents(documents, query)

        # Step 2: Synthesize themes across documents
        theme_results = self.synthesize_themes(document_results, query)

        # Return complete results
        return {
            "document_results": document_results,
            "themes": theme_results.get("themes", ""),
            "status": "success" if theme_results.get("status") == "success" else "partial",
            "query": query,
            "model_info": self.config
        }