"""
SealVera LangChain callback handler.

Logs every chain, LLM call, tool use, and agent action to SealVera automatically.

Usage:
    from sealvera.callbacks import SealVeraCallbackHandler
    import sealvera

    sealvera.init(api_key="sv_...", agent="my-langchain-agent")
    handler = SealVeraCallbackHandler()

    # With ChatOpenAI
    from langchain_openai import ChatOpenAI
    llm = ChatOpenAI(callbacks=[handler])

    # With a chain
    chain = LLMChain(llm=llm, prompt=prompt, callbacks=[handler])

    # Globally for all chains in a run
    chain.invoke({"input": "..."}, config={"callbacks": [handler]})
"""

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

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.outputs import LLMResult
    _LANGCHAIN_AVAILABLE = True
except ImportError:
    try:
        from langchain.callbacks.base import BaseCallbackHandler
        from langchain.schema import LLMResult
        _LANGCHAIN_AVAILABLE = True
    except ImportError:
        _LANGCHAIN_AVAILABLE = False
        # Stub so the import doesn't fail
        class BaseCallbackHandler:
            pass


class SealVeraCallbackHandler(BaseCallbackHandler if _LANGCHAIN_AVAILABLE else object):
    """
    LangChain callback handler that logs decisions to SealVera.

    Captures:
    - LLM calls (on_llm_start / on_llm_end / on_llm_error)
    - Chain runs (on_chain_start / on_chain_end / on_chain_error)
    - Tool calls (on_tool_start / on_tool_end / on_tool_error)
    - Agent actions (on_agent_action / on_agent_finish)
    """

    def __init__(self, client=None, agent: str = None, log_llm: bool = True,
                 log_chains: bool = True, log_tools: bool = True):
        """
        Args:
            client:     SealVera client (uses default if omitted)
            agent:      Agent name override
            log_llm:    Log raw LLM calls (default True)
            log_chains: Log chain runs (default True)
            log_tools:  Log tool calls (default True)
        """
        if _LANGCHAIN_AVAILABLE:
            super().__init__()

        self._client    = client
        self._agent     = agent
        self._log_llm   = log_llm
        self._log_chains = log_chains
        self._log_tools  = log_tools
        self._timers: Dict[str, float] = {}
        self._inputs: Dict[str, Any]   = {}

    def _get_client(self):
        if self._client:
            return self._client
        from . import _default_client
        return _default_client

    def _timer_start(self, run_id) -> None:
        self._timers[str(run_id)] = time.time()

    def _timer_end(self, run_id) -> float:
        return round(time.time() - self._timers.pop(str(run_id), time.time()), 3)

    # ── LLM ──────────────────────────────────────────────────────────────────

    def on_llm_start(self, serialized: Dict, prompts: List[str], run_id=None, **kwargs):
        if not self._log_llm:
            return
        self._timer_start(run_id)
        self._inputs[str(run_id)] = {"prompts": [p[:500] for p in (prompts or [])]}

    def on_llm_end(self, response: Any, run_id=None, **kwargs):
        if not self._log_llm:
            return
        elapsed = self._timer_end(run_id)
        input_data = self._inputs.pop(str(run_id), {})

        # Extract token usage if available
        usage = {}
        if hasattr(response, "llm_output") and response.llm_output:
            token_usage = response.llm_output.get("token_usage", {})
            if token_usage:
                usage = {
                    "prompt_tokens":     token_usage.get("prompt_tokens"),
                    "completion_tokens": token_usage.get("completion_tokens"),
                    "total_tokens":      token_usage.get("total_tokens"),
                }

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

        # Extract text output
        output_text = ""
        if hasattr(response, "generations") and response.generations:
            try:
                output_text = response.generations[0][0].text[:500]
            except Exception:
                pass

        self._get_client().log(
            action="llm_call",
            decision="COMPLETED",
            agent=self._agent,
            model=model,
            input=input_data,
            output={"text": output_text, "latency_s": elapsed, **usage},
            reasoning=[{"factor": "latency_s", "value": str(elapsed),
                        "signal": "safe", "explanation": "LLM call completed successfully"}],
        )

    def on_llm_error(self, error: Exception, run_id=None, **kwargs):
        if not self._log_llm:
            return
        elapsed = self._timer_end(run_id)
        self._get_client().log(
            action="llm_call",
            decision="ERROR",
            agent=self._agent,
            input=self._inputs.pop(str(run_id), {}),
            output={"error": str(error)[:300], "latency_s": elapsed},
            reasoning=[{"factor": "exception", "value": type(error).__name__,
                        "signal": "block", "explanation": str(error)[:200]}],
        )

    # ── Chain ─────────────────────────────────────────────────────────────────

    def on_chain_start(self, serialized: Dict, inputs: Dict, run_id=None, **kwargs):
        if not self._log_chains:
            return
        self._timer_start(run_id)
        chain_name = (serialized or {}).get("name", "chain")
        self._inputs[f"chain:{run_id}"] = {
            "chain": chain_name,
            "inputs": {k: str(v)[:200] for k, v in (inputs or {}).items()},
        }

    def on_chain_end(self, outputs: Dict, run_id=None, **kwargs):
        if not self._log_chains:
            return
        elapsed    = self._timer_end(run_id)
        input_data = self._inputs.pop(f"chain:{run_id}", {})
        self._get_client().log(
            action="chain_run",
            decision="COMPLETED",
            agent=self._agent,
            input=input_data,
            output={"outputs": {k: str(v)[:300] for k, v in (outputs or {}).items()},
                    "latency_s": elapsed},
            reasoning=[{"factor": "latency_s", "value": str(elapsed),
                        "signal": "safe", "explanation": "Chain completed successfully"}],
        )

    def on_chain_error(self, error: Exception, run_id=None, **kwargs):
        if not self._log_chains:
            return
        elapsed    = self._timer_end(run_id)
        input_data = self._inputs.pop(f"chain:{run_id}", {})
        self._get_client().log(
            action="chain_run",
            decision="FAILED",
            agent=self._agent,
            input=input_data,
            output={"error": str(error)[:300], "latency_s": elapsed},
            reasoning=[{"factor": "exception", "value": type(error).__name__,
                        "signal": "block", "explanation": str(error)[:200]}],
        )

    # ── Tool ─────────────────────────────────────────────────────────────────

    def on_tool_start(self, serialized: Dict, input_str: str, run_id=None, **kwargs):
        if not self._log_tools:
            return
        self._timer_start(run_id)
        tool_name = (serialized or {}).get("name", "tool")
        self._inputs[f"tool:{run_id}"] = {"tool": tool_name, "input": input_str[:300]}

    def on_tool_end(self, output: str, run_id=None, **kwargs):
        if not self._log_tools:
            return
        elapsed    = self._timer_end(run_id)
        input_data = self._inputs.pop(f"tool:{run_id}", {})
        self._get_client().log(
            action=f"tool_call",
            decision="COMPLETED",
            agent=self._agent,
            input=input_data,
            output={"result": str(output)[:500], "latency_s": elapsed},
            reasoning=[{"factor": "tool", "value": input_data.get("tool", "unknown"),
                        "signal": "safe", "explanation": "Tool executed successfully"}],
        )

    def on_tool_error(self, error: Exception, run_id=None, **kwargs):
        if not self._log_tools:
            return
        elapsed    = self._timer_end(run_id)
        input_data = self._inputs.pop(f"tool:{run_id}", {})
        self._get_client().log(
            action="tool_call",
            decision="ERROR",
            agent=self._agent,
            input=input_data,
            output={"error": str(error)[:300], "latency_s": elapsed},
            reasoning=[{"factor": "exception", "value": type(error).__name__,
                        "signal": "block", "explanation": str(error)[:200]}],
        )

    # ── Agent ─────────────────────────────────────────────────────────────────

    def on_agent_action(self, action: Any, run_id=None, **kwargs):
        tool  = getattr(action, "tool",       "unknown")
        input_ = getattr(action, "tool_input", {})
        self._get_client().log(
            action="agent_action",
            decision="RESPONDED",
            agent=self._agent,
            input={"tool": tool, "tool_input": str(input_)[:300]},
            output={"log": getattr(action, "log", "")[:300]},
            reasoning=[{"factor": "tool_selected", "value": tool,
                        "signal": "safe", "explanation": "Agent chose tool to execute"}],
        )

    def on_agent_finish(self, finish: Any, run_id=None, **kwargs):
        self._get_client().log(
            action="agent_finish",
            decision="COMPLETED",
            agent=self._agent,
            input={},
            output={"output": str(getattr(finish, "return_values", {}))[:500],
                    "log": getattr(finish, "log", "")[:300]},
            reasoning=[{"factor": "outcome", "value": "finished",
                        "signal": "safe", "explanation": "Agent completed its run"}],
        )
