"""ChatWZRD — LangChain ChatModel with WZRD velocity-aware model selection.

Wraps langchain_openai.ChatOpenAI. On every invocation, calls wzrd.pick()
to select the best model based on live adoption velocity, then delegates
to ChatOpenAI with that model. Routing happens per call, not at init,
so the model tracks momentum in real time.

Usage:
    from langchain_wzrd import ChatWZRD

    llm = ChatWZRD(task="code", openai_api_key="sk-or-...")
    response = llm.invoke("Write a Python sort function")

With auto-reporting for rewards:
    llm = ChatWZRD(
        task="code",
        openai_api_key="sk-or-...",
        wzrd_report=True,
        wzrd_keypair="~/.config/solana/id.json",
    )
"""

from __future__ import annotations

import logging
import threading
from typing import Any, Optional, Sequence

from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import BaseMessage
from langchain_core.outputs import ChatResult
from langchain_openai import ChatOpenAI

import wzrd

logger = logging.getLogger("langchain_wzrd")


class ChatWZRD(BaseChatModel):
    """LangChain ChatModel that routes to the best model via WZRD velocity signals.

    Every call to invoke/generate selects a model using wzrd.pick(), which
    reads live adoption velocity from the WZRD signal API. The selected model
    is served through OpenRouter via ChatOpenAI.

    Args:
        task: Task hint for model selection — "code", "chat", "reasoning", "general".
        openai_api_key: OpenRouter API key (passed as openai_api_key to ChatOpenAI).
        openai_api_base: API base URL. Defaults to OpenRouter.
        fallback: Model to use if WZRD signal is unavailable.
        candidates: Optional allowlist of model names to consider.
        wzrd_report: If True, report each routing decision for rewards.
        wzrd_keypair: Solana keypair path for agent auth (required if wzrd_report=True).
        temperature: Passed through to the underlying ChatOpenAI.
        max_tokens: Passed through to the underlying ChatOpenAI.
    """

    task: str = "general"
    openai_api_key: str = ""
    openai_api_base: str = "https://openrouter.ai/api/v1"
    fallback: str = "meta-llama/llama-3.3-70b-instruct"
    candidates: Optional[list[str]] = None
    wzrd_report: bool = False
    wzrd_keypair: Optional[str] = None
    temperature: float = 0.7
    max_tokens: Optional[int] = None

    _agent: Any = None

    model_config = {"arbitrary_types_allowed": True}

    def model_post_init(self, __context: Any) -> None:
        if self.wzrd_report and self._agent is None:
            self._agent = self._make_agent()

    @property
    def _llm_type(self) -> str:
        return "wzrd"

    def _generate(
        self,
        messages: list[BaseMessage],
        stop: Optional[list[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        choice = wzrd.pick_details(
            self.task,
            candidates=self.candidates,
            fallback=self.fallback,
        )

        model_name = choice.model
        logger.info("wzrd: %s → %s (trend=%s, confidence=%s)", self.task, model_name, choice.trend, choice.confidence)

        llm = ChatOpenAI(
            model=model_name,
            openai_api_key=self.openai_api_key,
            openai_api_base=self.openai_api_base,
            temperature=self.temperature,
            max_tokens=self.max_tokens,
        )

        result = llm._generate(messages, stop=stop, run_manager=run_manager, **kwargs)

        if self.wzrd_report and self._agent is not None:
            thread = threading.Thread(
                target=self._report_choice,
                args=(choice,),
                daemon=True,
            )
            thread.start()

        return result

    @property
    def _identifying_params(self) -> dict[str, Any]:
        return {
            "task": self.task,
            "fallback": self.fallback,
            "candidates": self.candidates,
            "wzrd_report": self.wzrd_report,
        }

    def _report_choice(self, choice: wzrd.PickResult) -> None:
        try:
            self._agent.report_pick(choice)
        except Exception:
            logger.warning("wzrd: auto-report failed for %s", choice.model, exc_info=True)

    def _make_agent(self) -> Any:
        try:
            agent = wzrd.WZRDAgent.from_env()
            agent.authenticate(self.wzrd_keypair)
            return agent
        except Exception:
            logger.warning("wzrd: agent auth failed; reporting disabled", exc_info=True)
            return None
