"""LangChain-backed AI provider implementations."""

from __future__ import annotations

import importlib
import uuid
from collections.abc import Callable, Iterable, Iterator
from contextlib import AbstractContextManager, nullcontext
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    SupportsInt,
    cast,
)

import structlog

from cleveragents.agents.plan_generation import PlanGenerationGraph
from cleveragents.domain.models.core import Change, Context, Plan, Project
from cleveragents.domain.providers.ai_provider import (
    ActorInvocationContext,
    AIProviderInterface,
    ProviderResponse,
)

GetOpenAICallback = Callable[[], AbstractContextManager[Any]]
_openai_callback: GetOpenAICallback | None = None

if _openai_callback is None:  # LangChain optional callback utilities
    try:
        callbacks_module = importlib.import_module("langchain.callbacks")
    except ModuleNotFoundError:  # pragma: no cover - optional dependency fallback
        pass
    else:  # pragma: no cover - exercised in provider tests via patching
        callback = getattr(callbacks_module, "get_openai_callback", None)
        if callable(callback):
            _openai_callback = cast(GetOpenAICallback, callback)

if TYPE_CHECKING:  # pragma: no cover - imported for type checking only
    from langchain_core.language_models import BaseLanguageModel

TokenEstimatorResult = int | float | SupportsInt | str | None
TokenEstimator = Callable[[str], TokenEstimatorResult]


class LangChainChatProvider(AIProviderInterface):
    """AI provider that uses a LangChain chat model with PlanGenerationGraph."""

    _DEFAULT_PROGRESS_MAP: ClassVar[dict[str, int]] = {
        "load_context": 15,
        "analyze_requirements": 40,
        "generate_plan": 70,
        "validate": 90,
    }

    def __init__(
        self,
        *,
        name: str,
        model_id: str,
        llm_factory: Callable[[str], BaseLanguageModel],
        max_retries: int = 3,
        supports_streaming: bool = True,
        progress_map: dict[str, int] | None = None,
    ) -> None:
        self._name = name
        self._model_id = model_id
        self._llm_factory = llm_factory
        self._max_retries = max(1, max_retries)
        self._supports_streaming = supports_streaming
        self._progress_map = progress_map or self._DEFAULT_PROGRESS_MAP.copy()
        self._logger = structlog.get_logger(__name__).bind(
            provider=name,
            model=model_id,
        )

    @property
    def name(self) -> str:  # pragma: no cover - simple accessor
        return self._name

    @name.setter
    def name(self, value: str) -> None:
        self._name = value

    @property
    def model_id(self) -> str:  # pragma: no cover - simple accessor
        return self._model_id

    @model_id.setter
    def model_id(self, value: str) -> None:
        self._model_id = value

    def generate_changes(
        self,
        project: Project,
        plan: Plan,
        contexts: list[Context],
        actor_context: ActorInvocationContext | None = None,
        progress_callback: Callable[[int], None] | None = None,
    ) -> ProviderResponse:
        """Generate code changes by running the LangGraph workflow."""

        llm = self._llm_factory(self._model_id)
        graph = PlanGenerationGraph(llm=llm, max_retries=self._max_retries)
        thread_id = f"provider-{uuid.uuid4()}"
        usage_tracker: Any | None = None

        def _run_workflow() -> dict[str, Any]:
            if progress_callback:
                progress_callback(5)

            if progress_callback and self._supports_streaming:
                return self._execute_with_streaming(
                    graph,
                    project,
                    plan,
                    contexts,
                    thread_id,
                    actor_context,
                    progress_callback,
                )

            state = graph.invoke(
                project,
                plan,
                contexts,
                thread_id=thread_id,
                actor_context=actor_context,
            )
            if progress_callback:
                progress_callback(90)
                progress_callback(100)
            return state

        try:
            with self._usage_tracker(llm) as tracker:
                usage_tracker = cast(Any | None, tracker)
                state = _run_workflow()
        except Exception as exc:  # pragma: no cover - defensive path
            if progress_callback:
                progress_callback(100)

            token_count = self._resolve_token_count(llm, plan, contexts, usage_tracker)
            cost = self._resolve_token_cost(usage_tracker)
            self._log_usage(model=self._model_id, tokens=token_count, cost=cost)
            return ProviderResponse(
                changes=[],
                model_used=self._model_id,
                token_count=token_count,
                error_message=str(exc),
            )

        token_count = self._resolve_token_count(llm, plan, contexts, usage_tracker)
        cost = self._resolve_token_cost(usage_tracker)
        self._log_usage(model=self._model_id, tokens=token_count, cost=cost)
        return self._build_response(state, token_count)

    def stream_changes(
        self,
        project: Project,
        plan: Plan,
        contexts: list[Context],
        actor_context: ActorInvocationContext | None = None,
        progress_callback: Callable[[int], None] | None = None,
    ) -> Iterator[dict[str, object]]:
        """Stream LangGraph events while generating code changes."""

        if not self._supports_streaming:

            def _fallback() -> Iterator[dict[str, object]]:
                response = self.generate_changes(
                    project,
                    plan,
                    contexts,
                    actor_context=actor_context,
                    progress_callback=progress_callback,
                )
                for node in (
                    "load_context",
                    "analyze_requirements",
                    "generate_plan",
                    "validate",
                ):
                    payload: dict[str, object] = {}
                    if node == "generate_plan":
                        payload = {
                            "status": "completed",
                            "change_count": len(response.changes),
                        }
                    elif node == "validate":
                        payload = {"status": "completed"}
                    yield {node: payload}
                yield {"__end__": {"response": response}}

            return _fallback()

        llm = self._llm_factory(self._model_id)
        graph = PlanGenerationGraph(llm=llm, max_retries=self._max_retries)
        thread_id = f"provider-{uuid.uuid4()}"

        def _stream() -> Iterator[dict[str, object]]:
            usage_tracker: Any | None = None
            state: dict[str, Any] = {
                "generated_changes": [],
                "validation_result": {},
                "error": None,
            }
            try:
                with self._usage_tracker(llm) as tracker:
                    usage_tracker = cast(Any | None, tracker)
                    if progress_callback:
                        progress_callback(5)
                    for event in graph.stream(
                        project,
                        plan,
                        contexts,
                        thread_id=thread_id,
                        actor_context=actor_context,
                    ):
                        node_name, payload = self._extract_event(event)
                        if isinstance(payload, dict):
                            state.update(cast(dict[str, Any], payload))
                        if node_name in {"__start__", "__end__"}:
                            continue
                        if progress_callback:
                            self._emit_progress(node_name, progress_callback)
                        yield {node_name: payload or {}}
                if progress_callback:
                    progress_callback(100)
            except Exception:
                if progress_callback:
                    progress_callback(100)
                raise

            token_count = self._resolve_token_count(llm, plan, contexts, usage_tracker)
            cost = self._resolve_token_cost(usage_tracker)
            self._log_usage(model=self._model_id, tokens=token_count, cost=cost)
            response = self._build_response(state, token_count)
            yield {"__end__": {"response": response}}

        return _stream()

    def _execute_with_streaming(
        self,
        graph: PlanGenerationGraph,
        project: Project,
        plan: Plan,
        contexts: list[Context],
        thread_id: str,
        actor_context: ActorInvocationContext | None,
        progress_callback: Callable[[int], None],
    ) -> dict[str, Any]:
        state: dict[str, Any] = {
            "generated_changes": [],
            "validation_result": {},
            "error": None,
        }
        try:
            for event in graph.stream(
                project,
                plan,
                contexts,
                thread_id=thread_id,
                actor_context=actor_context,
            ):
                node_name, payload = self._extract_event(event)
                self._emit_progress(node_name, progress_callback)
                if isinstance(payload, dict):
                    state.update(cast(dict[str, Any], payload))
            progress_callback(100)
        except Exception:
            progress_callback(100)
            raise
        return state

    def _usage_tracker(self, llm: BaseLanguageModel) -> AbstractContextManager[Any]:
        if _openai_callback is not None and self._is_openai_llm(llm):
            return _openai_callback()
        return nullcontext()

    def _is_openai_llm(self, llm: BaseLanguageModel) -> bool:
        module = getattr(llm.__class__, "__module__", "")
        return module.startswith("langchain_openai")

    def _resolve_token_count(
        self,
        llm: BaseLanguageModel,
        plan: Plan,
        contexts: list[Context],
        usage_tracker: Any | None,
    ) -> int:
        if usage_tracker is not None:
            tokens = getattr(usage_tracker, "total_tokens", None)
            if tokens is not None:
                try:
                    return int(tokens)
                except (TypeError, ValueError):
                    pass
        return self._estimate_token_usage(llm, plan, contexts)

    def _resolve_token_cost(self, usage_tracker: Any | None) -> float | None:
        if usage_tracker is None:
            return None
        cost = getattr(usage_tracker, "total_cost", None)
        if cost is None:
            return None
        try:
            return float(cost)
        except (TypeError, ValueError):  # pragma: no cover - defensive fallback
            return None

    def _log_usage(
        self,
        *,
        model: str,
        tokens: int | None,
        cost: float | None,
    ) -> None:
        if tokens is None and cost is None:
            return
        self._logger.info(
            "provider.token_usage",
            model=model,
            tokens=tokens,
            cost=cost,
        )

    def _build_response(
        self,
        state: dict[str, Any],
        token_count: int,
    ) -> ProviderResponse:
        validation = self._safe_validation_result(state)
        error_message = self._safe_error_message(state)
        if not error_message and validation:
            status = str(validation.get("status", "")).upper()
            if status == "FAIL":
                error_message = validation.get("message", "Validation failed")

        generated_changes = self._safe_generated_changes(state)

        return ProviderResponse(
            changes=generated_changes,
            model_used=self._model_id,
            token_count=token_count,
            error_message=error_message,
        )

    def _extract_event(self, event: dict[str, Any]) -> tuple[str, Any]:
        try:
            key, payload = next(iter(event.items()))
        except StopIteration:  # pragma: no cover - defensive guard
            return "__unknown__", {}
        return key, payload

    def _emit_progress(
        self, node_name: str, progress_callback: Callable[[int], None]
    ) -> None:
        percent = self._progress_map.get(node_name)
        if percent is not None:
            progress_callback(percent)

    def _estimate_token_usage(
        self,
        llm: BaseLanguageModel,
        plan: Plan,
        contexts: list[Context],
    ) -> int:
        get_tokens = getattr(llm, "get_num_tokens", None)
        if not callable(get_tokens):
            return 0

        estimator = cast(TokenEstimator, get_tokens)

        prompt_text = plan.prompt or ""
        for ctx in contexts:
            content = getattr(ctx, "content", None) or ""
            if content:
                prompt_text += f"\n{content[:2000]}"

        try:
            tokens = estimator(prompt_text)
        except Exception:  # pragma: no cover - defensive fallback
            return 0

        normalized_tokens = tokens
        if normalized_tokens is None:
            return 0

        try:
            return int(normalized_tokens)
        except (TypeError, ValueError):  # pragma: no cover - defensive fallback
            return 0

    def _safe_generated_changes(self, state: dict[str, Any]) -> list[Change]:
        value = state.get("generated_changes")
        if isinstance(value, list):
            return cast(list[Change], value)
        return []

    def _safe_validation_result(self, state: dict[str, Any]) -> dict[str, Any]:
        value = state.get("validation_result")
        if isinstance(value, dict):
            return cast(dict[str, Any], value)
        return {}

    def _safe_error_message(self, state: dict[str, Any]) -> str | None:
        value = state.get("error")
        if value is None:
            return None
        extracted = self._extract_error_message(value)
        if extracted:
            return extracted
        return str(value)

    def _extract_error_message(self, value: Any, depth: int = 0) -> str | None:
        if value is None:
            return None
        if depth > 10:
            return str(value)
        if isinstance(value, BaseException):
            return str(value)
        if isinstance(value, str):
            return value
        if isinstance(value, dict):
            typed_value: dict[str, Any] = cast(dict[str, Any], value)
            # Prefer nested error containers before top-level summaries
            priority_keys = (
                "error",
                "errors",
                "retry_error",
                "retry",
                "last_attempt",
                "exception",
                "exc",
                "cause",
                "detail",
                "reason",
                "message",
            )
            for key in priority_keys:
                if key in typed_value:
                    nested = self._extract_error_message(typed_value[key], depth + 1)
                    if nested:
                        return nested
            # Fallback to searching remaining values
            for nested_value in cast(Iterable[Any], typed_value.values()):
                nested = self._extract_error_message(nested_value, depth + 1)
                if nested:
                    return nested
            return None
        if isinstance(value, (list, tuple, set)):
            for item in cast(Iterable[Any], value):
                nested = self._extract_error_message(item, depth + 1)
                if nested:
                    return nested
            return None
        return str(value)
