#!/usr/bin/env python3
"""
Example of using a custom LLM with a custom retriever in Local Deep Research.

This demonstrates how to integrate your own LLM implementation and custom
retrieval system for programmatic access.
"""

from typing import List, Dict
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_core.retrievers import Document
from langchain_community.vectorstores import FAISS

# Import the search system
from local_deep_research.search_system import AdvancedSearchSystem

# Re-enable logging after import
from loguru import logger
import sys

logger.remove()
# diagnose=False: loguru defaults to True, which renders repr() of every
# local in every traceback frame on exception. Users copy this snippet
# into their own scripts, so leaving the default on would propagate the
# credential-in-traceback leak (#4185) wherever the snippet lands.
logger.add(
    sys.stderr,
    level="INFO",
    format="{time} {level} {message}",
    diagnose=False,
)
logger.enable("local_deep_research")


class CustomRetriever:
    """Custom retriever that can fetch from multiple sources."""

    def __init__(self):
        # Initialize with sample documents for demonstration
        self.documents = [
            {
                "content": "Quantum computing uses quantum bits (qubits) that can exist in superposition, "
                "allowing parallel computation of multiple states simultaneously.",
                "title": "Quantum Computing Fundamentals",
                "source": "quantum_basics.pdf",
                "metadata": {"topic": "quantum", "year": 2024},
            },
            {
                "content": "Machine learning algorithms can be categorized into supervised, unsupervised, "
                "and reinforcement learning approaches, each suited for different tasks.",
                "title": "ML Algorithm Categories",
                "source": "ml_overview.pdf",
                "metadata": {"topic": "ml", "year": 2024},
            },
            {
                "content": "Neural networks are inspired by biological neurons and consist of interconnected "
                "nodes that process information through weighted connections.",
                "title": "Neural Network Architecture",
                "source": "nn_architecture.pdf",
                "metadata": {"topic": "neural_networks", "year": 2023},
            },
            {
                "content": "Natural language processing enables computers to understand, interpret, and "
                "generate human language, powering applications like chatbots and translation.",
                "title": "NLP Applications",
                "source": "nlp_apps.pdf",
                "metadata": {"topic": "nlp", "year": 2024},
            },
        ]

        # Create embeddings for similarity search
        logger.info("Initializing custom retriever with embeddings...")
        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")

        # Create vector store from documents
        docs = [
            Document(
                page_content=doc["content"],
                metadata={
                    "title": doc["title"],
                    "source": doc["source"],
                    **doc["metadata"],
                },
            )
            for doc in self.documents
        ]
        self.vectorstore = FAISS.from_documents(docs, self.embeddings)

    def retrieve(self, query: str, k: int = 3) -> List[Dict]:
        """Custom retrieval logic."""
        logger.info(f"Custom Retriever: Searching for '{query}'")

        # Use vector similarity search
        similar_docs = self.vectorstore.similarity_search(query, k=k)

        # Convert to expected format
        results = []
        for i, doc in enumerate(similar_docs):
            results.append(
                {
                    "title": doc.metadata.get("title", f"Document {i + 1}"),
                    "link": doc.metadata.get("source", "custom_source"),
                    "snippet": doc.page_content[:150] + "...",
                    "full_content": doc.page_content,
                    "rank": i + 1,
                    "metadata": doc.metadata,
                }
            )

        logger.info(
            f"Custom Retriever: Found {len(results)} relevant documents"
        )
        return results


class CustomSearchEngine:
    """Adapter to integrate custom retriever with the search system."""

    def __init__(self, retriever: CustomRetriever, settings_snapshot=None):
        self.retriever = retriever
        self.settings_snapshot = settings_snapshot or {}

    def run(self, query: str, research_context=None) -> List[Dict]:
        """Execute search using custom retriever."""
        return self.retriever.retrieve(query, k=5)


def main():
    """Demonstrate custom LLM and retriever integration."""
    print("=== Custom LLM and Retriever Example ===\n")

    # 1. Create custom LLM (just using regular Ollama for simplicity)
    print("1. Initializing LLM...")
    llm = ChatOllama(model="gemma3:12b", temperature=0.7)

    # 2. Create custom retriever
    print("2. Setting up custom retriever...")
    custom_retriever = CustomRetriever()

    # 3. Create settings
    settings = {
        "search.iterations": 2,
        "search.questions_per_iteration": 3,
        "search.strategy": "source-based",
        "rate_limiting.enabled": False,  # Disable rate limiting for custom setup
    }

    # 4. Create search engine adapter
    print("3. Creating search engine adapter...")
    search_engine = CustomSearchEngine(custom_retriever, settings)

    # 5. Initialize the search system
    print("4. Initializing AdvancedSearchSystem with custom components...")
    # Pass programmatic_mode=True to avoid database dependencies
    search_system = AdvancedSearchSystem(
        llm=llm,
        search=search_engine,
        settings_snapshot=settings,
        programmatic_mode=True,
    )

    # 6. Run research queries
    queries = [
        "How do quantum computers differ from classical computers?",
        "What are the main types of machine learning algorithms?",
    ]

    for query in queries:
        print(f"\n{'=' * 60}")
        print(f"Research Query: {query}")
        print("=" * 60)

        result = search_system.analyze_topic(query)

        # Display results
        print("\n=== FINDINGS ===")
        print(result["formatted_findings"])

        # Show metadata
        print("\n=== SEARCH METADATA ===")
        print(f"• Total findings: {len(result['findings'])}")
        print(f"• Iterations: {result['iterations']}")

        # Get actual sources from all_links_of_system or search_results
        all_links = result.get("all_links_of_system", [])
        for finding in result.get("findings", []):
            if "search_results" in finding and finding["search_results"]:
                all_links = finding["search_results"]
                break

        print(f"• Sources found: {len(all_links)}")
        if all_links and len(all_links) > 0:
            print("\n=== SOURCES ===")
            for i, link in enumerate(all_links[:5], 1):  # Show first 5
                if isinstance(link, dict):
                    title = link.get("title", "No title")
                    url = link.get("link", link.get("source", "Unknown"))
                    print(f"  [{i}] {title}")
                    print(f"      URL: {url}")

        # Show generated questions
        if result.get("questions_by_iteration"):
            print("\n=== RESEARCH QUESTIONS GENERATED ===")
            for iteration, questions in result[
                "questions_by_iteration"
            ].items():
                print(f"\nIteration {iteration}:")
                for q in questions[:3]:  # Show first 3 questions
                    print(f"  • {q}")

    print("\n✓ Custom LLM and Retriever integration successful!")


if __name__ == "__main__":
    main()