"""LangInsight callback handler implementation.

TypeScript @langinsight/langchain と同じ仕様:
- metadata (userId, sessionId, modelName?, ...) でオプションを渡す
- on_llm_end で input と output の2本のトレースを並列送信
- result で送信完了後の input/output トレースを取得
- score(trace_id, value) でフィードバック送信
"""

import logging
from concurrent.futures import Future
from datetime import datetime, timezone
from threading import Thread
from typing import Any, Callable, Dict, Literal, Optional
from uuid import UUID

import httpx
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult

logger = logging.getLogger(__name__)

# API レスポンスの Trace 型（camelCase で返る）
TraceDict = Dict[str, Any]


class CallbackHandler(BaseCallbackHandler):
    """LangInsight callback handler for LangChain.

    TypeScript の LangInsight.CallbackHandler と同様に:
    - 各 LLM 終了時に input と output の2本のトレースを /traces/external に並列送信
    - result で送信完了後の { input, output } を取得可能（Future）
    - score(trace_id, value) で /feedbacks/external にフィードバック送信
    """

    name = "langinsight_callback_handler"

    def __init__(
        self,
        api_key: str,
        endpoint: str,
        metadata: Dict[str, Any],
        on_success: Optional[Callable[[TraceDict, TraceDict], None]] = None,
        on_failure: Optional[Callable[[Exception], None]] = None,
    ):
        """Initialize the LangInsight callback handler.

        Args:
            api_key: LangInsight API key.
            endpoint: LangInsight API endpoint (e.g. https://api.langinsight.example.com).
            metadata: Must include "userId" and "sessionId". Optional "modelName".
                      Other keys are sent as meta on traces.
            on_success: Optional callback (input_trace, output_trace) when both traces are sent.
            on_failure: Optional callback (error) when trace submission fails.
        """
        super().__init__()
        self._api_key = api_key
        self._endpoint = endpoint.rstrip("/")
        self._metadata = dict(metadata)
        self._on_success = on_success
        self._on_failure = on_failure
        self._run_prompts: Dict[UUID, list[str]] = {}
        self._result_future: Future[tuple[TraceDict, TraceDict]] = Future()

    def _headers(self) -> Dict[str, str]:
        return {"x-api-key": self._api_key, "Content-Type": "application/json"}

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: list[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Store prompts for this run so we can send them as input trace in on_llm_end."""
        self._run_prompts[run_id] = list(prompts)

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Send input and output traces to LangInsight in parallel (same as TypeScript)."""
        if not response.generations or not response.generations[0]:
            return

        generation = response.generations[0][0]
        message = getattr(generation, "message", None)
        if not message:
            return

        # Timing
        run_start = getattr(response, "llm_output", None)
        start_time = (
            run_start.get("start_time")
            if isinstance(run_start, dict)
            else kwargs.get("start_time")
        )
        end_time = (
            run_start.get("end_time")
            if isinstance(run_start, dict)
            else kwargs.get("end_time")
        )
        started_at = (
            datetime.fromtimestamp(start_time, tz=timezone.utc).isoformat()
            if start_time is not None
            else datetime.now(timezone.utc).isoformat()
        )
        ended_at = (
            datetime.fromtimestamp(end_time, tz=timezone.utc).isoformat()
            if end_time is not None
            else datetime.now(timezone.utc).isoformat()
        )

        # Input content: stored prompt from on_llm_start
        prompts = self._run_prompts.pop(run_id, [])
        input_content = prompts[-1] if prompts else ""

        # Output
        text = generation.text
        response_metadata = getattr(message, "response_metadata", None) or {}
        usage_metadata = getattr(message, "usage_metadata", None) or {}
        if isinstance(usage_metadata, dict):
            input_tokens = usage_metadata.get("input_tokens", 0) or usage_metadata.get("total_tokens", 0)
            output_tokens = usage_metadata.get("output_tokens", 0)
        else:
            input_tokens = output_tokens = 0

        user_id = self._metadata.get("userId") or self._metadata.get("user_id", "")
        session_id = self._metadata.get("sessionId") or self._metadata.get("session_id", "")
        model_name = self._metadata.get("modelName") or self._metadata.get("model_name")
        model = (
            model_name
            or response_metadata.get("model")
            or response_metadata.get("model_name")
            or "unknown"
        )
        meta = {
            k: v
            for k, v in self._metadata.items()
            if k not in ("userId", "user_id", "sessionId", "session_id", "modelName", "model_name")
        }

        def post_trace(payload: Dict[str, Any]) -> TraceDict:
            with httpx.Client(timeout=30.0) as client:
                r = client.post(
                    f"{self._endpoint}/traces/external",
                    headers=self._headers(),
                    json=payload,
                )
                r.raise_for_status()
                return r.json()

        input_payload = {
            "userId": user_id,
            "sessionId": session_id,
            "model": model,
            "token": input_tokens,
            "content": input_content,
            "type": "input",
            "meta": meta,
            "startedAt": started_at,
            "endedAt": ended_at,
        }
        output_payload = {
            "userId": user_id,
            "sessionId": session_id,
            "model": model,
            "token": output_tokens,
            "content": text,
            "type": "output",
            "meta": meta,
            "startedAt": started_at,
            "endedAt": ended_at,
        }

        results: list[Optional[TraceDict]] = [None, None]
        errors: list[Optional[Exception]] = [None, None]

        def do_input() -> None:
            try:
                results[0] = post_trace(input_payload)
            except Exception as e:
                errors[0] = e

        def do_output() -> None:
            try:
                results[1] = post_trace(output_payload)
            except Exception as e:
                errors[1] = e

        t1 = Thread(target=do_input)
        t2 = Thread(target=do_output)
        t1.start()
        t2.start()
        t1.join()
        t2.join()

        err = errors[0] or errors[1]
        if err:
            logger.exception("Failed to send traces: %s", err)
            if self._on_failure:
                self._on_failure(err if isinstance(err, Exception) else Exception(str(err)))
            self._result_future.set_exception(err)
            return

        input_trace = results[0]
        output_trace = results[1]
        if input_trace is None or output_trace is None:
            ex = Exception("One or both trace requests failed")
            self._result_future.set_exception(ex)
            if self._on_failure:
                self._on_failure(ex)
            return
        self._result_future.set_result((input_trace, output_trace))
        if self._on_success:
            self._on_success(input_trace, output_trace)

    @property
    def result(self) -> Future[tuple[TraceDict, TraceDict]]:
        """Future that resolves to (input_trace, output_trace) when both traces are sent.

        Use: input_trace, output_trace = handler.result.result(timeout=10)
        or with concurrent.futures: input_trace, output_trace = handler.result.result()
        """
        return self._result_future

    def score(self, trace_id: str, value: Literal[1, -1]) -> None:
        """Send feedback (thumbs up/down) for the given trace.

        trace_id should be the output trace id from result (e.g. output_trace["id"]).
        """
        user_id = self._metadata.get("userId") or self._metadata.get("user_id", "")
        session_id = self._metadata.get("sessionId") or self._metadata.get("session_id", "")
        with httpx.Client(timeout=30.0) as client:
            r = client.post(
                f"{self._endpoint}/feedbacks/external",
                headers=self._headers(),
                json={
                    "traceId": trace_id,
                    "userId": user_id,
                    "sessionId": session_id,
                    "value": value,
                },
            )
            if not r.is_success:
                logger.error("Failed to send feedback: %s %s", r.status_code, r.text)
