"""PydanticAI adapter for the DeepSeek reward-tuning advisor."""

from __future__ import annotations

import hashlib
import json
import os
from dataclasses import dataclass
from typing import Any

from .schema import validate_proposal

SYSTEM_PROMPT = """你是 Unitree Go2 强化学习奖励调参专家。
只根据提供的数值配置、训练曲线摘要和固定评估结果提出下一轮稀疏修改。
必须优先保持当前任务的客观安全门槛：Flat速度跟踪/跌倒，Obstacle无碰撞到达/跌倒。
每轮最多修改四个白名单标量，不得改变符号、函数、传感器、评估协议或结构。
不要建议 Python 代码、命令、文件路径或白名单外参数。输出必须符合 RewardProposal schema。
"""


class AdvisorUnavailable(RuntimeError):
    pass


@dataclass(frozen=True)
class AdvisorConfig:
    api_key: str | None
    base_url: str = "https://api.deepseek.com"
    model: str = "deepseek-v4-flash"

    @classmethod
    def from_environment(cls) -> AdvisorConfig:
        return cls(
            api_key=os.environ.get("DEEPSEEK_API_KEY"),
            base_url=os.environ.get("MUJOCO_TUNING_AGENT_BASE_URL", "https://api.deepseek.com"),
            model=os.environ.get("MUJOCO_TUNING_AGENT_MODEL", "deepseek-v4-flash"),
        )


class DeepSeekAdvisor:
    def __init__(self, config: AdvisorConfig | None = None):
        self.config = config or AdvisorConfig.from_environment()
        self._cached_agent = None

    def capability(self) -> dict[str, Any]:
        try:
            import pydantic_ai  # noqa: F401
        except ImportError:
            installed = False
        else:
            installed = True
        return {
            "configured": bool(self.config.api_key) and installed,
            "apiKeyConfigured": bool(self.config.api_key),
            "frameworkInstalled": installed,
            "model": self.config.model,
            "baseUrl": self.config.base_url,
        }

    def _agent(self):
        if self._cached_agent is not None:
            return self._cached_agent
        if not self.config.api_key:
            raise AdvisorUnavailable("未配置 DEEPSEEK_API_KEY")
        try:
            import httpx2
            from pydantic import BaseModel, Field
            from pydantic_ai import Agent, PromptedOutput
            from pydantic_ai.models.openai import OpenAIChatModel
            from pydantic_ai.providers.openai import OpenAIProvider
        except ImportError as error:
            raise AdvisorUnavailable(
                "缺少 PydanticAI，请安装 training_server/requirements.txt"
            ) from error

        class RewardProposalOutput(BaseModel):
            weights: dict[str, float] = Field(default_factory=dict)
            params: dict[str, float] = Field(default_factory=dict)
            rationale: str = Field(min_length=1, max_length=2000)
            expected_impact: dict[str, str] = Field(default_factory=dict)
            confidence: float = Field(ge=0.0, le=1.0)

        proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("ALL_PROXY")
        if proxy and proxy.startswith("socks://"):
            proxy = "socks5://" + proxy.removeprefix("socks://")
        http_client = httpx2.AsyncClient(proxy=proxy, trust_env=False, timeout=60.0)
        provider = OpenAIProvider(
            base_url=self.config.base_url, api_key=self.config.api_key, http_client=http_client
        )
        model = OpenAIChatModel(self.config.model, provider=provider)  # type: ignore[arg-type]
        self._cached_agent = Agent(
            model,
            output_type=PromptedOutput(RewardProposalOutput),
            system_prompt=SYSTEM_PROMPT,
            retries=2,
            model_settings={"temperature": 0.2},
        )
        return self._cached_agent

    def propose(self, context: dict[str, Any], previous: dict) -> dict[str, Any]:
        prompt = json.dumps(context, ensure_ascii=False, separators=(",", ":"), allow_nan=False)
        result = self._agent().run_sync(prompt)
        output = result.output
        patch = validate_proposal(
            {"weights": dict(output.weights), "params": dict(output.params)},
            previous,
            task_id=context.get("task", "Unitree-Go2-Flat"),
        )
        try:
            usage = result.usage()
            usage_value = {
                key: getattr(usage, key)
                for key in ("requests", "input_tokens", "output_tokens", "total_tokens")
                if getattr(usage, key, None) is not None
            }
        except (AttributeError, TypeError):
            usage_value = {}
        return {
            "patch": patch,
            "rationale": output.rationale,
            "expectedImpact": dict(output.expected_impact),
            "confidence": float(output.confidence),
            "promptHash": hashlib.sha256(prompt.encode()).hexdigest(),
            "usage": usage_value,
            "model": self.config.model,
        }

    def test_connection(self) -> dict[str, Any]:
        base = {
            "weights": {"track_linear_velocity": 1.0},
            "params": {},
            "instruction": "仅返回一个合法示例：把 track_linear_velocity 改为 1.1。",
        }
        # A minimal full previous config is supplied by callers for actual proposals;
        # connectivity probing only verifies the provider and structured response path.
        result = self._agent().run_sync(json.dumps(base, ensure_ascii=False))
        return {"ok": True, "model": self.config.model, "outputType": type(result.output).__name__}