"""
LangChain integration for RAGLint.
Provides a callback handler to automatically trace LangChain executions.
"""

from typing import Any, Dict, List, Optional, Union
from uuid import UUID

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.outputs import LLMResult
except ImportError:
    # Create dummy class if langchain is not installed
    class BaseCallbackHandler:
        pass
    class LLMResult:
        pass

from raglint.instrumentation import Monitor

class RAGLintCallbackHandler(BaseCallbackHandler):
    """
    Callback handler for LangChain to automatically log events to RAGLint.
    
    Usage:
        from raglint.integrations.langchain import RAGLintCallbackHandler
        
        chain.invoke(input, config={"callbacks": [RAGLintCallbackHandler()]})
    """

    def __init__(self):
        self.monitor = Monitor()

    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
    ) -> Any:
        """Run when chain starts running."""
        self.monitor.log_event("chain_start", {
            "trace_id": str(run_id),
            "parent_id": str(parent_run_id) if parent_run_id else None,
            "inputs": inputs,
            "name": serialized.get("name", "LangChain")
        })

    def on_chain_end(
        self, outputs: Dict[str, Any], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
    ) -> Any:
        """Run when chain ends running."""
        self.monitor.log_event("chain_end", {
            "trace_id": str(run_id),
            "outputs": outputs
        })

    def on_chain_error(
        self, error: Union[Exception, KeyboardInterrupt], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
    ) -> Any:
        """Run when chain errors."""
        self.monitor.log_event("chain_error", {
            "trace_id": str(run_id),
            "error": str(error)
        })

    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
    ) -> Any:
        """Run when LLM starts running."""
        self.monitor.log_event("llm_start", {
            "trace_id": str(run_id),
            "parent_id": str(parent_run_id) if parent_run_id else None,
            "prompts": prompts,
            "model": serialized.get("kwargs", {}).get("model_name")
        })

    def on_llm_end(
        self, response: LLMResult, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
    ) -> Any:
        """Run when LLM ends running."""
        self.monitor.log_event("llm_end", {
            "trace_id": str(run_id),
            "generations": [[g.text for g in gen] for gen in response.generations]
        })

    def on_retriever_start(
        self, serialized: Dict[str, Any], query: str, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
    ) -> Any:
        """Run when Retriever starts running."""
        self.monitor.log_event("retriever_start", {
            "trace_id": str(run_id),
            "parent_id": str(parent_run_id) if parent_run_id else None,
            "query": query
        })

    def on_retriever_end(
        self, documents: Any, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
    ) -> Any:
        """Run when Retriever ends running."""
        # Handle both list of Documents and other formats
        docs_info = []
        try:
            for doc in documents:
                docs_info.append({
                    "content": getattr(doc, "page_content", str(doc))[:200] + "...",
                    "metadata": getattr(doc, "metadata", {})
                })
        except Exception:
            docs_info = str(documents)

        self.monitor.log_event("retriever_end", {
            "trace_id": str(run_id),
            "documents": docs_info
        })
