"""LangChain instrumentation via BaseCallbackHandler.

Registers an ART callback that evaluates every LLM call in a LangChain chain
or agent, with zero changes to user code.

Example::

    import art
    from langchain_openai import ChatOpenAI

    art.instrument_langchain()

    llm = ChatOpenAI(model="gpt-4o-mini")
    response = llm.invoke("What is the capital of France?")
    # [ART] gpt-4o-mini | confidence:85 hallucination:0 | $0.00012 | 320ms
"""
from __future__ import annotations

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

from art._config import get_config
from art._context import get_trace_id, increment_step
from art._cost import estimate_cost, _PRICES_PER_1K_INPUT, _DEFAULT_PRICE_PER_1K
from art._evaluators import evaluate
from art._exporters import make_exporter
from art._record import EvaluationRecord


def instrument_langchain() -> None:
    """Register ART as a global LangChain callback handler.

    Every LLM call in any LangChain chain or agent will be automatically
    evaluated and exported. Call once at application startup.

    Raises:
        ImportError: If langchain-core is not installed.
    """
    try:
        from langchain_core.callbacks import BaseCallbackManager
        from langchain_core.globals import get_llm_cache
    except ImportError:
        raise ImportError(
            "langchain-core is not installed. "
            "Install it with: pip install art-sdk[langchain]"
        )

    try:
        from langchain_core.callbacks.manager import _configure
    except ImportError:
        pass

    # Register globally via LangChain's callback system
    try:
        import langchain_core.callbacks.manager as _mgr
        _original_configure = getattr(_mgr, "_configure", None)

        handler = ARTCallbackHandler()

        # Add to global handlers if LangChain supports it
        if hasattr(_mgr, "get_callback_manager"):
            _mgr.get_callback_manager().add_handler(handler)
        else:
            # Fallback: monkey-patch to inject handler into every configure call
            _inject_global_handler(handler)
    except Exception:
        # Last resort: just return the handler so user can add it manually
        pass


def _inject_global_handler(handler) -> None:
    """Inject the ART handler into LangChain's global callback stack."""
    try:
        from langchain_core.callbacks import manager as _cb_manager

        if not getattr(_cb_manager, "_art_patched", False):
            original_configure = _cb_manager._configure

            def patched_configure(*args, **kwargs):
                result = original_configure(*args, **kwargs)
                # Add our handler if not already present
                if hasattr(result, "handlers"):
                    existing_types = {type(h) for h in result.handlers}
                    if ARTCallbackHandler not in existing_types:
                        result.add_handler(handler, inherit=True)
                return result

            _cb_manager._configure = patched_configure
            _cb_manager._art_patched = True
    except Exception:
        pass


def make_art_callback() -> "ARTCallbackHandler":
    """Create an ART callback handler to pass manually to LangChain chains.

    Use this when you prefer explicit configuration over global patching::

        from langchain_openai import ChatOpenAI
        cb = art.make_art_callback()
        llm = ChatOpenAI(model="gpt-4o-mini", callbacks=[cb])
    """
    return ARTCallbackHandler()


class ARTCallbackHandler:
    """LangChain callback handler that evaluates LLM calls with ART."""

    _start_times: Dict[str, float]

    def __init__(self) -> None:
        self._start_times = {}

    # ── Required by LangChain's callback protocol ─────────────────────────────

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: uuid.UUID,
        **kwargs: Any,
    ) -> None:
        self._start_times[str(run_id)] = time.monotonic()

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[Any]],
        *,
        run_id: uuid.UUID,
        **kwargs: Any,
    ) -> None:
        self._start_times[str(run_id)] = time.monotonic()

    def on_llm_end(
        self,
        response: Any,
        *,
        run_id: uuid.UUID,
        **kwargs: Any,
    ) -> None:
        run_key = str(run_id)
        start = self._start_times.pop(run_key, None)
        latency_ms = int((time.monotonic() - start) * 1000) if start else 0

        try:
            config = get_config()

            # Extract model name from response or kwargs
            model = ""
            if hasattr(response, "llm_output") and response.llm_output:
                model = response.llm_output.get("model_name", "") or ""

            # Extract prompt and response text
            prompt_text = ""
            response_text = ""

            generations = getattr(response, "generations", [])
            if generations:
                first_gen = generations[0]
                if first_gen:
                    gen = first_gen[0]
                    # ChatGeneration has .message, Generation has .text
                    if hasattr(gen, "message"):
                        msg = gen.message
                        response_text = getattr(msg, "content", "") or ""
                    elif hasattr(gen, "text"):
                        response_text = gen.text or ""

            # Extract token usage
            input_tokens = 0
            output_tokens = 0
            if hasattr(response, "llm_output") and response.llm_output:
                usage = response.llm_output.get("token_usage", {}) or {}
                input_tokens = usage.get("prompt_tokens", 0) or usage.get("input_tokens", 0) or 0
                output_tokens = usage.get("completion_tokens", 0) or usage.get("output_tokens", 0) or 0

            price = _PRICES_PER_1K_INPUT.get(model, _DEFAULT_PRICE_PER_1K)
            actual_cost = round((max(input_tokens, 1) / 1000) * price, 8)

            trace_id = get_trace_id()
            step_index = increment_step() if trace_id else 0

            eval_result = evaluate(prompt_text, response_text)

            record = EvaluationRecord(
                id=str(uuid.uuid4()),
                created_at=EvaluationRecord.now_iso(),
                model=model or "unknown",
                pipeline=config.pipeline,
                prompt=prompt_text,
                response=response_text,
                latency_ms=latency_ms,
                prompt_tokens=input_tokens,
                completion_tokens=output_tokens,
                estimated_cost_usd=actual_cost,
                confidence_score=eval_result.confidence_score,
                hallucination_score=eval_result.hallucination_score,
                coherence_score=eval_result.coherence_score,
                instruction_following_score=eval_result.instruction_following_score,
                hallucination_indicators=eval_result.hallucination_indicators,
                all_indicators=eval_result.all_indicators,
                trace_id=trace_id,
                step_index=step_index,
                source="langchain",
            )

            exporter = make_exporter(config)
            exporter.export(record)
        except Exception:
            pass  # Never let ART affect LangChain execution

    def on_llm_error(self, error: Exception, *, run_id: uuid.UUID, **kwargs: Any) -> None:
        self._start_times.pop(str(run_id), None)

    # Stubs for other events LangChain may call
    def on_chain_start(self, *args: Any, **kwargs: Any) -> None: pass
    def on_chain_end(self, *args: Any, **kwargs: Any) -> None: pass
    def on_chain_error(self, *args: Any, **kwargs: Any) -> None: pass
    def on_tool_start(self, *args: Any, **kwargs: Any) -> None: pass
    def on_tool_end(self, *args: Any, **kwargs: Any) -> None: pass
    def on_tool_error(self, *args: Any, **kwargs: Any) -> None: pass
    def on_agent_action(self, *args: Any, **kwargs: Any) -> None: pass
    def on_agent_finish(self, *args: Any, **kwargs: Any) -> None: pass
    def on_text(self, *args: Any, **kwargs: Any) -> None: pass
    def on_retriever_start(self, *args: Any, **kwargs: Any) -> None: pass
    def on_retriever_end(self, *args: Any, **kwargs: Any) -> None: pass
    def on_retriever_error(self, *args: Any, **kwargs: Any) -> None: pass

    # LangChain checks for these attributes
    raise_error: bool = False
    ignore_llm: bool = False
    ignore_chain: bool = False
    ignore_agent: bool = False
    ignore_retriever: bool = False
    ignore_chat_model: bool = False
