#!/usr/bin/env python3
"""
LLM Config — Reads agent_models.json + llm_providers.json to resolve
role → provider → model → API details.

agent_models.json:  maps agent role → model id (format: "provider/model-name")
llm_providers.json:  maps provider → baseUrl, apiKey, api type, model list
"""

import json
import logging
import os
from pathlib import Path
import time
from typing import Any, Optional

ROOT = Path(__file__).parent.parent.parent

_BUILTIN_AGENT_DEFAULTS: dict[str, dict[str, Any]] = {
    "AI 策略开发师": {
        "role": "AI 策略开发师",
        "role_key": "ai_strategy_developer",
        "description": "专用于策略定制、文章导入、多链接聚合与策略草稿编译。",
        "model": "deepseek/deepseek-v4-pro",
        "status": "ready",
    },
    "趋势分析师": {
        "role": "趋势分析师",
        "role_key": "trend_analyst",
        "description": "宏观方向指挥官。基于多时段技术数据 + 账户状态决策 LONG/SHORT/NEUTRAL。",
        "model": "deepseek/deepseek-v4-pro",
        "status": "ready",
    },
    "风险管理员": {
        "role": "风险管理员",
        "role_key": "risk_manager",
        "description": "信号二次验证。每个信号发出前做方向合理性检查、置信度校准、风险标记。",
        "model": "deepseek/deepseek-v4-pro",
        "status": "ready",
    },
    "M2 决策主脑": {
        "role": "M2 决策主脑",
        "role_key": "m2_brain",
        "description": "多智能体聚合决策主脑。",
        "model": "deepseek/deepseek-v4-flash",
        "status": "ready",
    },
    "本地巡查(兜底)": {
        "role": "本地巡查(兜底)",
        "role_key": "local_patrol",
        "description": "本地轻量兜底巡查模型。",
        "model": "omlx/Qwen3.5-4B-MLX-4bit",
        "status": "ready",
    },
    "离线训练师": {
        "role": "离线训练师",
        "role_key": "trainer",
        "description": "离线训练、复盘与参数建议模型。",
        "model": "deepseek/deepseek-v4-pro",
        "status": "ready",
    },
}


def _safe_load_json(path: Path, fallback_corrupt: bool = False) -> dict | None:
    """Load JSON with corruption detection + .bak recovery."""
    import logging
    _log = logging.getLogger("llm.config")
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError) as e:
        _log.error("Corrupted config file: %s — %s", path, e)
        # Try .bak recovery
        bak = path.with_suffix(path.suffix + ".bak")
        if bak.exists():
            try:
                _log.warning("Recovering from backup: %s", bak)
                return json.loads(bak.read_text(encoding="utf-8"))
            except Exception:
                pass
        # Rename corrupted file out of the way
        try:
            corrupt = path.with_suffix(path.suffix + f".corrupted_{int(time.time())}")
            path.rename(corrupt)
            _log.warning("Renamed corrupted file to %s", corrupt)
        except Exception:
            pass
        return {} if fallback_corrupt else None


def _get_config_dir() -> Path:
    try:
        import config_manager

        return config_manager.get_config_dir()
    except Exception:
        fallback = ROOT / "config"
        fallback.mkdir(exist_ok=True)
        return fallback


def _get_llm_providers_path() -> Optional[Path]:
    """Return path to the trading system's own LLM provider config (llm_providers.json)."""
    try:
        import config_manager
        own_config = config_manager.get_config_dir() / "llm_providers.json"
        if own_config.exists():
            return own_config
    except Exception:
        pass
    # Last resort: try source tree config
    for candidate in [
        ROOT / "llm_providers.json",
    ]:
        if candidate.exists():
            return candidate
    return None


def _get_custom_providers_path() -> Path:
    return _get_config_dir() / "custom_providers.json"


def load_agent_models() -> dict[str, dict[str, Any]]:
    """Return {role: {role, model, status, ...}} from personas + UI overrides.

    Each persona file contains role, role_key, description, model, temperature,
    max_tokens, timeout_sec, system_prompt, output_schema.
    Persona files are treated as defaults. agent_models.json can override the
    bound model/status for the same role and may append extra roles.
    """
    personas_dir = _get_config_dir() / "personas"
    _log = logging.getLogger("llm.config")
    agents: dict[str, dict[str, Any]] = {}
    if personas_dir.is_dir():
        for pf in sorted(personas_dir.glob("*.json")):
            try:
                persona = json.loads(pf.read_text(encoding="utf-8"))
                role = persona.get("role", "")
                if role:
                    persona["status"] = "ready"
                    agents[role] = persona
            except Exception:
                _log.warning("Failed to load persona file: %s", pf)
    if not agents:
        agents = {role: dict(config) for role, config in _BUILTIN_AGENT_DEFAULTS.items()}

    # Legacy/UI overrides: role -> selected model
    config_file = _get_config_dir() / "agent_models.json"
    if config_file.exists():
        try:
            data = _safe_load_json(config_file, fallback_corrupt=True)
            if data and isinstance(data, list):
                for item in data:
                    if not isinstance(item, dict):
                        continue
                    role = str(item.get("role", "") or "").strip()
                    if not role:
                        continue
                    merged = dict(agents.get(role, _BUILTIN_AGENT_DEFAULTS.get(role, {"role": role})))
                    merged["role"] = role
                    if item.get("model"):
                        merged["model"] = str(item.get("model", "")).strip()
                    if item.get("status"):
                        merged["status"] = item.get("status")
                    if item.get("temperature") is not None:
                        merged["temperature"] = item.get("temperature")
                    if item.get("max_tokens"):
                        merged["max_tokens"] = item.get("max_tokens")
                    if item.get("timeout_sec") is not None:
                        merged["timeout_sec"] = item.get("timeout_sec")
                    if item.get("description") and not merged.get("description"):
                        merged["description"] = item.get("description")
                    if item.get("api_key_configured") is not None:
                        merged["api_key_configured"] = item["api_key_configured"]
                    agents[role] = merged
        except Exception:
            pass
    return agents


def _deep_merge_provider(original: dict, override: dict) -> dict:
    """Deep-merge override into original.

    - Simple fields (baseUrl, api, apiKey, etc.): override wins if non-empty
    - models list: merge by 'id' — custom models replace system models with same id,
      and new custom models are appended. System models not mentioned in custom are kept.
    """
    merged = dict(original)
    for key, val in override.items():
        if key == "_name":
            continue
        if key == "models" and isinstance(val, list) and isinstance(merged.get("models"), list):
            # Merge models by id: custom takes priority, system kept if not overridden
            sys_models = {m["id"]: m for m in merged["models"] if isinstance(m, dict) and "id" in m}
            for m in val:
                if isinstance(m, dict) and "id" in m:
                    sys_models[m["id"]] = m  # custom overrides or appends
            merged["models"] = list(sys_models.values())
        elif key == "apiKey":
            # Only override API key if override provides a non-trivial, non-placeholder value
            if val and val.strip() and val != "sk-test" and "YOUR_" not in val.upper():
                merged["apiKey"] = val
            # Otherwise keep original (which may be empty — env var fallback happens later)
        elif isinstance(val, (str, int, float, bool)) and val:
            # Non-empty scalar: override wins
            merged[key] = val
        elif val:
            # Non-empty complex value: override wins
            merged[key] = val
    return merged


def load_llm_providers() -> dict[str, dict[str, Any]]:
    """Return merged provider config from llm_providers.json (or legacy openclaw.json).

    Merge strategy:
      1. Load system providers from llm_providers.json (authoritative for baseUrl/api/models)
      2. Load custom providers from custom_providers.json
      3. Deep-merge: custom fields override system fields per-provider
      4. Models are merged by 'id' — custom can add or replace specific models
      5. API keys from custom are accepted ONLY if non-trivial (not sk-test/placeholder)
    """
    providers: dict[str, dict[str, Any]] = {}
    _log = logging.getLogger("llm.config")
    config_path = _get_llm_providers_path()
    # Fallback to legacy openclaw.json when llm_providers.json is absent
    if not config_path or not config_path.exists():
        try:
            import config_manager
            legacy = config_manager.get_llm_config_dir() / "openclaw.json"
            if legacy.exists():
                config_path = legacy
        except Exception:
            pass
    if config_path and config_path.exists():
        cfg = _safe_load_json(config_path, fallback_corrupt=True)
        if cfg:
            # Support two formats:
            #   New (llm_providers.json): {"providers": {"deepseek": {...}}}
            #   Legacy (openclaw.json):   {"models": {"providers": {"deepseek": {...}}}}
            if "providers" in cfg and isinstance(cfg["providers"], dict):
                system_providers = cfg["providers"]
            else:
                system_providers = cfg.get("models", {}).get("providers", {})
            if isinstance(system_providers, dict):
                for name, info in system_providers.items():
                    if isinstance(info, dict):
                        providers[name] = {**info, "_name": name}

    custom_path = _get_custom_providers_path()
    if custom_path.exists():
        custom_data = _safe_load_json(custom_path, fallback_corrupt=True) or {}
        custom_providers = custom_data.get("providers", {})
        if isinstance(custom_providers, dict):
            for name, info in custom_providers.items():
                    if not isinstance(info, dict):
                        continue
                    if name in providers:
                        # Deep-merge: custom overrides system, but preserves system models
                        providers[name] = _deep_merge_provider(providers[name], info)
                    else:
                        # Brand new custom provider
                        providers[name] = {**info, "_name": name}

        overrides = custom_data.get("overrides", {})
        if isinstance(overrides, dict):
            for name, info in overrides.items():
                if not isinstance(info, dict):
                    continue
                current = dict(providers.get(name, {"_name": name}))
                pricing = info.get("pricing")
                if isinstance(pricing, dict):
                    current["pricing"] = pricing
                providers[name] = current

    return providers


def resolve_agent_model(role: str) -> Optional[dict[str, str]]:
    """Get the {role, model, status} config for a given agent role."""
    agents = load_agent_models()
    return agents.get(role)


def resolve_provider_for_model(model_id: str) -> tuple[Optional[str], Optional[str], Optional[dict[str, Any]]]:
    """Parse "provider/model-name" and return (provider_name, model_name, provider_config).

    Returns (None, None, None) if provider not found.
    """
    if "/" not in model_id:
        return None, model_id, None

    provider_name, model_name = model_id.split("/", 1)
    providers = load_llm_providers()
    provider_config = providers.get(provider_name)
    return provider_name, model_name, provider_config


def get_api_key(provider_config: dict[str, Any]) -> str:
    """Extract API key from provider config, with env var fallback."""
    api_key = provider_config.get("apiKey", "")
    if api_key:
        return api_key

    # Try common env var patterns
    provider_name = provider_config.get("_name", "")
    env_map = {
        "deepseek": "DEEPSEEK_API_KEY",
        "openai": "OPENAI_API_KEY",
        "anthropic": "ANTHROPIC_API_KEY",
        "claude": "ANTHROPIC_API_KEY",
        "google": "GOOGLE_API_KEY",
        "gemini": "GOOGLE_API_KEY",
        "minimax": "MINIMAX_API_KEY",
    }
    for name, env_var in env_map.items():
        if name in str(provider_config.get("baseUrl", "")).lower() or provider_name.lower() == name:
            val = os.environ.get(env_var, "")
            if val:
                return val

    return ""
