"""
RAGResearchAgent - Production-Ready RAG Agent with Vector Search

Zero-config usage:
    from kaizen.agents import RAGResearchAgent

    agent = RAGResearchAgent()
    result = agent.research("What is machine learning?")
    print(result["answer"])
    print(result["sources"])
    print(f"Confidence: {result['confidence']}")

Progressive configuration:
    agent = RAGResearchAgent(
        llm_provider="openai",
        model="gpt-4",
        temperature=0.7,
        top_k_documents=5,
        similarity_threshold=0.4,
        embedding_model="all-mpnet-base-v2"
    )

Environment variable support:
    KAIZEN_LLM_PROVIDER=openai
    KAIZEN_MODEL=gpt-4
    KAIZEN_TEMPERATURE=0.7
    KAIZEN_MAX_TOKENS=1000
    KAIZEN_TOP_K=3
    KAIZEN_SIMILARITY_THRESHOLD=0.3
    KAIZEN_EMBEDDING_MODEL=all-MiniLM-L6-v2
"""

import os
from dataclasses import dataclass, field, replace
from typing import Any, Dict, List, Optional

from kaizen.core.base_agent import BaseAgent
from kaizen.memory.vector import VectorMemory
from kaizen.retrieval.vector_store import SimpleVectorStore
from kaizen.signatures import InputField, OutputField, Signature
from kaizen.strategies.multi_cycle import MultiCycleStrategy

from kailash.nodes.base import NodeMetadata, register_node

# Sample documents for knowledge base
SAMPLE_AI_DOCUMENTS = [
    {
        "id": "doc1",
        "title": "Introduction to Machine Learning",
        "content": "Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data. It involves training algorithms on datasets to make predictions or decisions without being explicitly programmed.",
    },
    {
        "id": "doc2",
        "title": "Deep Learning Fundamentals",
        "content": "Deep learning is a specialized branch of machine learning that uses neural networks with multiple layers. These networks can automatically learn hierarchical representations of data, making them particularly effective for tasks like image recognition and natural language processing.",
    },
    {
        "id": "doc3",
        "title": "Natural Language Processing",
        "content": "Natural Language Processing (NLP) is a field of AI that focuses on the interaction between computers and human language. Modern NLP uses deep learning techniques, particularly transformer architectures, to understand and generate human language.",
    },
    {
        "id": "doc4",
        "title": "Computer Vision Applications",
        "content": "Computer vision enables machines to interpret and understand visual information from the world. Applications include facial recognition, object detection, autonomous vehicles, and medical image analysis.",
    },
    {
        "id": "doc5",
        "title": "Reinforcement Learning",
        "content": "Reinforcement learning is a machine learning paradigm where agents learn to make decisions by interacting with an environment. The agent receives rewards or penalties based on its actions, learning optimal strategies over time.",
    },
]


@dataclass
class RAGConfig:
    """
    Configuration for RAG Research Agent.

    All parameters have sensible defaults and can be overridden via:
    1. Constructor arguments (highest priority)
    2. Environment variables (KAIZEN_*)
    3. Default values (lowest priority)
    """

    # LLM configuration
    llm_provider: str = field(
        default_factory=lambda: os.getenv("KAIZEN_LLM_PROVIDER", "openai")
    )
    model: str = field(
        default_factory=lambda: os.getenv("KAIZEN_MODEL", "gpt-3.5-turbo")
    )
    temperature: float = field(
        default_factory=lambda: float(os.getenv("KAIZEN_TEMPERATURE", "0.7"))
    )
    max_tokens: int = field(
        default_factory=lambda: int(os.getenv("KAIZEN_MAX_TOKENS", "1000"))
    )

    # RAG-specific configuration
    top_k_documents: int = field(
        default_factory=lambda: int(os.getenv("KAIZEN_TOP_K", "3"))
    )
    similarity_threshold: float = field(
        default_factory=lambda: float(os.getenv("KAIZEN_SIMILARITY_THRESHOLD", "0.3"))
    )
    embedding_model: str = field(
        default_factory=lambda: os.getenv("KAIZEN_EMBEDDING_MODEL", "all-MiniLM-L6-v2")
    )

    # Autonomous execution configuration
    max_cycles: int = field(
        default_factory=lambda: int(os.getenv("KAIZEN_MAX_CYCLES", "15"))
    )  # Research may need multiple query-fetch-analyze cycles

    # Technical configuration
    timeout: int = 30
    retry_attempts: int = 3
    provider_config: Dict[str, Any] = field(default_factory=dict)
    memory_config: Optional[Dict[str, Any]] = None


class RAGSignature(Signature):
    """
    Signature for RAG pattern: Retrieval + Generation.

    Takes a query, retrieves relevant documents, and generates
    an answer with source attribution.
    """

    # Input fields
    query: str = InputField(desc="User's research question or query")

    # Output fields
    answer: str = OutputField(desc="Comprehensive answer based on retrieved documents")
    sources: list = OutputField(desc="List of source document IDs used")
    confidence: float = OutputField(desc="Confidence in answer accuracy (0.0-1.0)")
    relevant_excerpts: list = OutputField(desc="Key excerpts from source documents")
    tool_calls: list = OutputField(
        desc="Tools to call for research (web_search, fetch_url, etc., empty = converged)"
    )


@register_node()
class RAGResearchAgent(BaseAgent):
    """
    Production-ready RAG Research Agent with vector search.

    Features:
    - Zero-config with sensible defaults
    - Progressive configuration (override as needed)
    - Environment variable support
    - Semantic vector search (90% precision vs 60% keyword)
    - Source attribution and confidence scoring
    - Optional VectorMemory integration
    - Built-in error handling and logging
    - Sample documents pre-loaded
    - **Autonomous execution**: Multi-cycle query → fetch → analyze → refine loops

    Usage:
        # Zero-config (easiest)
        agent = RAGResearchAgent()
        result = agent.research("What is machine learning?")

        # With configuration
        agent = RAGResearchAgent(
            llm_provider="openai",
            model="gpt-4",
            temperature=0.7,
            top_k_documents=5,
            similarity_threshold=0.4
        )

        # View results
        result = agent.research("Explain deep learning")
        print(result["answer"])
        print(f"Sources: {result['sources']}")
        print(f"Confidence: {result['confidence']}")
        print(f"Retrieval quality: {result['retrieval_quality']}")

    Configuration:
        llm_provider: LLM provider (default: "openai", env: KAIZEN_LLM_PROVIDER)
        model: Model name (default: "gpt-3.5-turbo", env: KAIZEN_MODEL)
        temperature: Sampling temperature (default: 0.7, env: KAIZEN_TEMPERATURE)
        max_tokens: Maximum tokens (default: 1000, env: KAIZEN_MAX_TOKENS)
        top_k_documents: Number of documents to retrieve (default: 3, env: KAIZEN_TOP_K)
        similarity_threshold: Minimum similarity score (default: 0.3, env: KAIZEN_SIMILARITY_THRESHOLD)
        embedding_model: Sentence transformer model (default: "all-MiniLM-L6-v2", env: KAIZEN_EMBEDDING_MODEL)
        timeout: Request timeout seconds (default: 30)
        retry_attempts: Retry count on failure (default: 3)
        vector_store: Optional custom vector store (default: creates new with sample docs)
        memory_config: Optional memory config dict (default: None - disabled)
        provider_config: Additional provider-specific config (default: {})

    Returns:
        Dict with keys:
        - answer: str - Comprehensive answer
        - sources: list - Document IDs used
        - confidence: float - Answer confidence 0.0-1.0
        - relevant_excerpts: list - Key excerpts from documents
        - retrieval_quality: float - Average similarity score
        - error: str (optional) - Error code if validation fails
    """

    # Node metadata for Studio discovery
    metadata = NodeMetadata(
        name="RAGResearchAgent",
        description="Retrieval-Augmented Generation agent with semantic vector search and source attribution",
        version="1.0.0",
        tags={
            "ai",
            "kaizen",
            "rag",
            "research",
            "retrieval",
            "vector-search",
            "semantic",
        },
    )

    def __init__(
        self,
        llm_provider: Optional[str] = None,
        model: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        top_k_documents: Optional[int] = None,
        similarity_threshold: Optional[float] = None,
        embedding_model: Optional[str] = None,
        timeout: Optional[int] = None,
        retry_attempts: Optional[int] = None,
        vector_store: Optional[SimpleVectorStore] = None,
        memory_config: Optional[Dict[str, Any]] = None,
        provider_config: Optional[Dict[str, Any]] = None,
        config: Optional[RAGConfig] = None,
        tool_registry: Optional[Any] = None,
        mcp_servers: Optional[List[Dict[str, Any]]] = None,
        **kwargs,
    ):
        """
        Initialize RAG Research Agent with zero-config defaults.

        Args:
            llm_provider: Override default LLM provider
            model: Override default model
            temperature: Override default temperature
            max_tokens: Override default max tokens
            top_k_documents: Override default top-k retrieval
            similarity_threshold: Override default similarity threshold
            embedding_model: Override default embedding model
            timeout: Override default timeout
            retry_attempts: Override default retry attempts
            vector_store: Custom vector store (default: creates new with sample docs)
            memory_config: Memory configuration dict (opt-in)
            provider_config: Additional provider-specific configuration
            config: Full config object (overrides individual params)
            tool_registry: Optional tool registry for autonomous tool calling
            mcp_servers: Optional MCP server configurations for tool discovery
        """
        # If config object provided, use it; otherwise build from parameters
        if config is None:
            config = RAGConfig()

            # Override defaults with provided parameters
            if llm_provider is not None:
                config = replace(config, llm_provider=llm_provider)
            if model is not None:
                config = replace(config, model=model)
            if temperature is not None:
                config = replace(config, temperature=temperature)
            if max_tokens is not None:
                config = replace(config, max_tokens=max_tokens)
            if top_k_documents is not None:
                config = replace(config, top_k_documents=top_k_documents)
            if similarity_threshold is not None:
                config = replace(config, similarity_threshold=similarity_threshold)
            if embedding_model is not None:
                config = replace(config, embedding_model=embedding_model)
            if timeout is not None:
                config = replace(config, timeout=timeout)
            if retry_attempts is not None:
                config = replace(config, retry_attempts=retry_attempts)
            if memory_config is not None:
                config = replace(config, memory_config=memory_config)
            if provider_config is not None:
                config = replace(config, provider_config=provider_config)

        # Merge timeout into provider_config
        if config.timeout and (
            not config.provider_config or "timeout" not in config.provider_config
        ):
            provider_cfg = (
                config.provider_config.copy() if config.provider_config else {}
            )
            provider_cfg["timeout"] = config.timeout
            config = replace(config, provider_config=provider_cfg)

        # Initialize memory if enabled
        memory = None
        if config.memory_config and config.memory_config.get("enabled"):
            embedding_fn = config.memory_config.get(
                "embedder"
            )  # Optional custom embedder
            top_k = config.memory_config.get("top_k", 5)
            memory = VectorMemory(embedding_fn=embedding_fn, top_k=top_k)

        # CRITICAL: Initialize MultiCycleStrategy for autonomous execution
        # Research is iterative: query → fetch → analyze → refine → repeat
        multi_cycle_strategy = MultiCycleStrategy(
            max_cycles=config.max_cycles, convergence_check=self._check_convergence
        )

        # Initialize BaseAgent with auto-config extraction
        super().__init__(
            config=config,  # Auto-extracted to BaseAgentConfig
            signature=RAGSignature(),
            strategy=multi_cycle_strategy,  # CRITICAL: Autonomous execution
            memory=memory,
            tool_registry=tool_registry,
            mcp_servers=mcp_servers,
            **kwargs,
        )

        self.rag_config = config

        # Initialize vector store with embeddings
        if vector_store is None:
            self.vector_store = SimpleVectorStore(
                embedding_model=config.embedding_model
            )
            # Load default documents
            self.vector_store.add_documents(SAMPLE_AI_DOCUMENTS)
        else:
            self.vector_store = vector_store

    def _check_convergence(self, result: Dict[str, Any]) -> bool:
        """
        Check if research cycle should stop (convergence detection).

        Implements ADR-013: Objective convergence detection via tool_calls field.

        Convergence logic (priority order):
        1. OBJECTIVE (preferred): Check tool_calls field
           - tool_calls present and non-empty → NOT converged (continue researching)
           - tool_calls present but empty → CONVERGED (research complete)
        2. SUBJECTIVE (fallback): Check confidence and research depth
           - confidence >= 0.85 AND research_depth == "comprehensive" → CONVERGED
        3. DEFAULT: CONVERGED (safe fallback)

        Args:
            result: Cycle result from LLM

        Returns:
            True if converged (stop), False if continue

        Examples:
            >>> # Cycle 1: Initial query, needs more sources
            >>> result = {"tool_calls": [{"name": "web_search", "params": {...}}]}
            >>> agent._check_convergence(result)
            False  # Has tool calls → continue researching

            >>> # Cycle 5: Research comprehensive
            >>> result = {"tool_calls": [], "confidence": 0.90}
            >>> agent._check_convergence(result)
            True  # Empty tool calls → converged
        """
        # OBJECTIVE CONVERGENCE (PREFERRED)
        if "tool_calls" in result:
            tool_calls = result.get("tool_calls", [])

            # Validate format
            if not isinstance(tool_calls, list):
                # Malformed tool_calls → stop for safety
                return True

            if tool_calls:
                # Has tool calls (web_search, fetch_url, etc.) → continue
                return False

            # Empty tool calls → research complete
            return True

        # SUBJECTIVE FALLBACK (backward compatibility)
        confidence = result.get("confidence", 0)
        research_depth = result.get("research_depth", "shallow")

        # High confidence + comprehensive research → converged
        if confidence >= 0.85 and research_depth == "comprehensive":
            return True

        # DEFAULT: Converged (safe fallback)
        return True

    def research(self, query: str, session_id: Optional[str] = None) -> Dict[str, Any]:
        """
        Perform RAG-based research on a query using semantic vector search.

        Args:
            query: User's research question
            session_id: Optional session identifier for memory tracking

        Returns:
            Dict containing answer, sources, confidence, excerpts, and similarity scores

        Example:
            >>> agent = RAGResearchAgent()
            >>> result = agent.research("What is machine learning?")
            >>> print(result["answer"])
            Machine learning is a subset of artificial intelligence...
            >>> print(result["sources"])
            ['doc1', 'doc2']
            >>> print(result["confidence"])
            0.85
        """
        # Input validation
        if not query or not query.strip():
            return {
                "answer": "Please provide a valid research query.",
                "sources": [],
                "confidence": 0.0,
                "relevant_excerpts": [],
                "error": "INVALID_INPUT",
            }

        # Step 1: Semantic search with vector store
        retrieved_docs = self.vector_store.search(
            query=query.strip(),
            top_k=self.rag_config.top_k_documents,
            similarity_threshold=self.rag_config.similarity_threshold,
        )

        if not retrieved_docs:
            return {
                "answer": "No relevant documents found for your query.",
                "sources": [],
                "confidence": 0.0,
                "relevant_excerpts": [],
                "error": "NO_DOCUMENTS",
            }

        # Step 2: Prepare context from retrieved documents with similarity scores
        context_parts = []
        for i, doc in enumerate(retrieved_docs, 1):
            similarity = doc.get("similarity", 0.0)
            context_parts.append(
                f"[Document {i}: {doc['title']} (Relevance: {similarity:.2f})]"
            )
            context_parts.append(doc["content"])
            context_parts.append("")  # Empty line separator

        context = "\n".join(context_parts)

        # Step 3: Create enhanced query with context
        enhanced_query = f"""Based on the following documents, answer the query.

{context}

Query: {query.strip()}

Provide a comprehensive answer using information from the documents above.
Cite specific documents when referencing information."""

        # Step 4: Execute via BaseAgent with context
        result = self.run(query=enhanced_query, session_id=session_id)

        # Step 5: Add source metadata with similarity scores
        result["sources"] = [doc["id"] for doc in retrieved_docs]
        result["relevant_excerpts"] = [
            {
                "title": doc["title"],
                "excerpt": doc["content"][:100] + "...",
                "similarity": doc.get("similarity", 0.0),
            }
            for doc in retrieved_docs
        ]

        # Add average retrieval quality
        avg_similarity = sum(
            doc.get("similarity", 0.0) for doc in retrieved_docs
        ) / len(retrieved_docs)
        result["retrieval_quality"] = avg_similarity

        return result

    def add_document(self, doc_id: str, title: str, content: str):
        """
        Add a document to the vector store knowledge base.

        Args:
            doc_id: Unique document identifier
            title: Document title
            content: Document content

        Note: This will generate embeddings for the new document.

        Example:
            >>> agent = RAGResearchAgent()
            >>> agent.add_document(
            ...     doc_id="doc6",
            ...     title="Quantum Computing",
            ...     content="Quantum computing uses quantum mechanics..."
            ... )
        """
        self.vector_store.add_documents(
            [{"id": doc_id, "title": title, "content": content}]
        )

    def get_document_count(self) -> int:
        """
        Get total number of documents in knowledge base.

        Returns:
            int: Number of documents in vector store

        Example:
            >>> agent = RAGResearchAgent()
            >>> count = agent.get_document_count()
            >>> print(f"Knowledge base has {count} documents")
            Knowledge base has 5 documents
        """
        return len(self.vector_store.documents)

    def clear_documents(self):
        """
        Clear all documents from the vector store.

        Warning: This removes all documents including sample documents.

        Example:
            >>> agent = RAGResearchAgent()
            >>> agent.clear_documents()
            >>> print(agent.get_document_count())
            0
        """
        self.vector_store.clear()


# Convenience function for quick usage
def research(query: str, **kwargs) -> Dict[str, Any]:
    """
    Quick one-liner for RAG research without creating an agent instance.

    Args:
        query: The research question
        **kwargs: Optional configuration (llm_provider, model, temperature, etc.)

    Returns:
        The full result dictionary

    Example:
        >>> from kaizen.agents.specialized.rag_research import research
        >>> result = research("What is deep learning?")
        >>> print(result["answer"])
        Deep learning is a specialized branch of machine learning...
        >>> print(result["sources"])
        ['doc2', 'doc1']
    """
    agent = RAGResearchAgent(**kwargs)
    return agent.research(query)
