"""
Advanced example of custom LLM integration with Local Deep Research.

This example demonstrates:
- Factory functions with configuration
- Error handling and retry logic
- Combining multiple LLMs
- Integration with custom retrievers
"""

import time
from typing import Any, Dict, List, Optional

from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from loguru import logger

from local_deep_research.api import (
    create_settings_snapshot,
    detailed_research,
    quick_summary,
)


class RetryLLM(BaseChatModel):
    """LLM wrapper that adds retry logic to any base LLM."""

    base_llm: BaseChatModel
    max_retries: int = 3
    retry_delay: float = 1.0

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Generate with retry logic."""
        last_error = None
        delay = self.retry_delay

        for attempt in range(self.max_retries):
            try:
                return self.base_llm._generate(
                    messages, stop, run_manager, **kwargs
                )
            except Exception as e:
                last_error = e
                if attempt < self.max_retries - 1:
                    logger.warning(
                        f"Attempt {attempt + 1} failed, retrying in {delay}s..."
                    )
                    time.sleep(delay)
                    delay *= 2  # Exponential backoff

        raise last_error

    @property
    def _llm_type(self) -> str:
        return f"retry_{self.base_llm._llm_type}"


class ConfigurableLLM(BaseChatModel):
    """LLM that can be configured with custom parameters."""

    model_name: str = "configurable-v1"
    response_style: str = "technical"
    max_length: int = 500
    include_confidence: bool = False

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Generate response based on configuration."""
        # Extract the query
        query = messages[-1].content if messages else "No query"

        # Build response based on style
        if self.response_style == "technical":
            response = (
                f"Technical Analysis ({self.model_name}): {query[:100]}..."
            )
        elif self.response_style == "simple":
            response = (
                f"Simple Answer: Based on the query about {query[:50]}..."
            )
        else:
            response = f"Response: Processing '{query[:50]}...'"

        # Limit length
        response = response[: self.max_length]

        # Add confidence if requested
        if self.include_confidence:
            response += "\n\nConfidence: High"  # Use descriptive confidence instead of hardcoded percentage

        message = AIMessage(content=response)
        generation = ChatGeneration(message=message)

        return ChatResult(generations=[generation])

    @property
    def _llm_type(self) -> str:
        return "configurable"


_DOMAIN_KNOWLEDGE: Dict[str, List[str]] = {
    "medical": ["diagnosis", "treatment", "symptoms", "medications"],
    "legal": ["contracts", "liability", "regulations", "compliance"],
    "technical": ["algorithms", "architecture", "performance", "scalability"],
    "finance": ["investments", "risk", "portfolio", "markets"],
}


class DomainExpertLLM(BaseChatModel):
    """LLM that specializes in specific domains."""

    domain: str = "general"
    expertise_level: float = 0.8

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Generate domain-specific response."""
        query = messages[-1].content if messages else ""

        # Check if query matches domain
        domain_terms = _DOMAIN_KNOWLEDGE.get(self.domain, [])
        relevance = sum(
            1 for term in domain_terms if term.lower() in query.lower()
        )

        if relevance > 0:
            response = f"[{self.domain.upper()} EXPERT - High Relevance]: "
        else:
            response = f"[{self.domain.upper()} EXPERT - General]: "

        response += f"Based on my {self.domain} expertise (level: {self.expertise_level}), "
        response += f"regarding '{query[:100]}...': This requires specialized knowledge."

        message = AIMessage(content=response)
        generation = ChatGeneration(message=message)

        return ChatResult(generations=[generation])

    @property
    def _llm_type(self) -> str:
        return f"expert_{self.domain}"


def create_configured_llm(config: Dict[str, Any]) -> BaseChatModel:
    """Factory function that creates LLMs based on configuration."""
    llm_type = config.get("type", "basic")

    if llm_type == "retry":
        # Create base LLM first
        base_config = config.get("base_config", {})
        base_llm = create_configured_llm(base_config)

        # Wrap with retry
        return RetryLLM(
            base_llm=base_llm,
            max_retries=config.get("max_retries", 3),
            retry_delay=config.get("retry_delay", 1.0),
        )

    if llm_type == "configurable":
        return ConfigurableLLM(
            model_name=config.get("model_name", "config-v1"),
            response_style=config.get("style", "technical"),
            max_length=config.get("max_length", 500),
            include_confidence=config.get("include_confidence", False),
        )

    if llm_type == "expert":
        return DomainExpertLLM(
            domain=config.get("domain", "general"),
            expertise_level=config.get("expertise_level", 0.8),
        )

    # Default fallback
    return ConfigurableLLM()


def main():
    logger.info("Advanced Custom LLM Integration Examples")
    logger.info("=" * 60)

    # Example 1: Using a retry wrapper
    logger.info("\n1. Retry Wrapper Example:")
    base_llm = ConfigurableLLM(response_style="technical")
    retry_llm = RetryLLM(base_llm=base_llm, max_retries=3)

    snapshot = create_settings_snapshot(
        provider="retry_tech",
        overrides={"search.tool": "wikipedia"},
    )
    result = quick_summary(
        query="Explain quantum computing applications",
        llms={"retry_tech": retry_llm},
        settings_snapshot=snapshot,
    )
    logger.info(f"Summary: {result['summary'][:200]}...")

    # Example 2: Multiple domain experts
    logger.info("\n\n2. Multiple Domain Experts:")
    experts = {
        "medical_expert": DomainExpertLLM(
            domain="medical", expertise_level=0.95
        ),
        "tech_expert": DomainExpertLLM(domain="technical", expertise_level=0.9),
        "finance_expert": DomainExpertLLM(
            domain="finance", expertise_level=0.85
        ),
    }

    # Medical query
    snapshot = create_settings_snapshot(
        provider="medical_expert",
        overrides={"search.tool": "pubmed"},
    )
    _ = quick_summary(
        query="What are the latest treatments for diabetes?",
        llms=experts,
        settings_snapshot=snapshot,
    )
    logger.info(
        "Medical summary retrieved successfully. Content not logged for privacy."
    )

    # Example 3: Factory with configuration
    logger.info("\n\n3. Factory Configuration Example:")

    # Configuration for a technical writer
    tech_writer_config = {
        "type": "configurable",
        "model_name": "tech-writer-v2",
        "style": "technical",
        "max_length": 1000,
        "include_confidence": True,
    }

    # Configuration for a retry wrapper around the technical writer
    robust_config = {
        "type": "retry",
        "max_retries": 5,
        "retry_delay": 0.5,
        "base_config": tech_writer_config,
    }

    snapshot = create_settings_snapshot(
        provider="robust_writer",
        overrides={"search.tool": "arxiv"},
    )
    result = quick_summary(
        query="How do neural networks learn?",
        llms={
            "robust_writer": lambda **kwargs: create_configured_llm(
                robust_config
            )
        },
        settings_snapshot=snapshot,
    )
    logger.info(f"Robust Writer: {result['summary'][:150]}...")

    # Example 4: Research pipeline with different LLMs
    logger.info("\n\n4. Multi-Stage Research Pipeline:")

    # Stage 1: Quick exploration with simple LLM
    simple_llm = ConfigurableLLM(response_style="simple", max_length=200)

    snapshot = create_settings_snapshot(provider="simple")
    initial = quick_summary(
        query="Climate change impacts on agriculture",
        llms={"simple": simple_llm},
        settings_snapshot=snapshot,
        iterations=1,
    )

    logger.info(f"Initial exploration: {initial['summary'][:100]}...")

    # Stage 2: Detailed research with expert
    expert_llm = DomainExpertLLM(domain="technical", expertise_level=0.95)

    snapshot = create_settings_snapshot(provider="expert")
    detailed = detailed_research(
        query="Climate change impacts on agriculture: focus on technology solutions",
        llms={"expert": expert_llm},
        settings_snapshot=snapshot,
        iterations=2,
    )

    logger.info(f"Expert analysis: {detailed['summary'][:150]}...")

    # Example 5: Combining custom LLMs with custom retrievers
    logger.info("\n\n5. Custom LLM + Retriever Combination:")

    # Mock retriever for demonstration
    class MockRetriever:
        def get_relevant_documents(self, query):
            return [
                {"page_content": f"Mock document about {query}", "metadata": {}}
            ]

    custom_llm = ConfigurableLLM(
        model_name="integrated-v1",
        response_style="technical",
        include_confidence=True,
    )

    snapshot = create_settings_snapshot(
        provider="integrated",
        overrides={"search.tool": "company_docs"},
    )
    result = quick_summary(
        query="Internal company policies on remote work",
        llms={"integrated": custom_llm},
        retrievers={"company_docs": MockRetriever()},
        settings_snapshot=snapshot,
    )

    logger.info(f"Integrated result: {result['summary'][:150]}...")


if __name__ == "__main__":
    main()