#!/usr/bin/env python3
"""
Agent Role Router — unified LLM entry point for all 5 agent roles.

Usage:
    from src.agent_role_router import call_llm, ROLE_TRAINER

    result = call_llm(ROLE_TRAINER, "Analyse this strategy performance...",
                      temperature=0.1, max_tokens=1500)

All roles → call_llm(role, prompt, **kwargs)
              │
              ├─ resolve_role_model(role) → deepseek/deepseek-v4-pro
              │
              ├─ _call_deepseek_api() → https://api.deepseek.com/v1/chat/completions
              ├─ _call_omlx_api()     → http://127.0.0.1:8000/v1/chat/completions
              └─ MiniMax / Anthropic  → routed via src.llm unified engine

Design: thin convenience wrapper around src.llm.call_agent_llm().
Returns raw text content (not a dict), raises on error.
"""

from __future__ import annotations

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

# ── Role constants ──────────────────────────────────────────────
from src.license.validator import sync_remote_secrets

ROLE_STRATEGY_DEVELOPER = "AI 策略开发师"
ROLE_TREND_ANALYST = "趋势分析师"
ROLE_RISK_MANAGER = "风险管理员"
ROLE_M2_BRAIN = "M2 决策主脑"
ROLE_LOCAL_PATROL = "本地巡查(兜底)"
ROLE_TRAINER = "离线训练师"

ALL_ROLES = [ROLE_STRATEGY_DEVELOPER, ROLE_TREND_ANALYST, ROLE_RISK_MANAGER, ROLE_M2_BRAIN, ROLE_LOCAL_PATROL, ROLE_TRAINER]

# ── Default system prompts per role ─────────────────────────────
from src.trading_context import get_trading_context as _get_ctx
_ctx = _get_ctx()

ROLE_SYSTEM_PROMPTS: dict[str, str] = {
    ROLE_STRATEGY_DEVELOPER: f"你是一个面向 {_ctx.role_identity_cn} 自动交易系统的 AI 策略开发师。你负责把自然语言需求、文章链接、市场上下文与策略目标编译成结构化、可治理、可回测的策略草稿 JSON 对象，并在需要时生成 Python 策略代码。只输出结构化结果，不做空泛解释。",
    ROLE_TREND_ANALYST: f"You are a strategic macro commander for a {_ctx.role_identity_en} algorithmic trading system. You make directional decisions based on technical data. Always output valid JSON.",
    ROLE_RISK_MANAGER: f"你是一个专业的{_ctx.role_identity_cn}交易信号验证与风险管理专家。只返回JSON，不做解释。",
    ROLE_M2_BRAIN: "你是一位专业的量化交易策略师，负责深度分析策略参数调整建议。",
    ROLE_LOCAL_PATROL: "You are a local patrol agent checking system health and trade validity.",
    ROLE_TRAINER: "你是一位专业的量化交易策略师与培训师，负责生成和改进交易策略代码。",
}


# ── Model resolution ────────────────────────────────────────────

def resolve_trainer_model() -> str:
    """Resolve the model for the trainer role.

    Priority: TRAINER_MODEL env var > agent_models.json (GUI config) > default deepseek-v4-pro
    """
    env_model = os.environ.get("TRAINER_MODEL", "")
    if env_model:
        return env_model

    try:
        import config_manager
        agents_file = config_manager.get_config_dir() / "agent_models.json"
        if agents_file.exists():
            agents = json.loads(agents_file.read_text(encoding="utf-8"))
            if isinstance(agents, list):
                for agent in agents:
                    if isinstance(agent, dict) and agent.get("role") == ROLE_TRAINER:
                        model = str(agent.get("model", "")).strip()
                        if model:
                            return model
    except Exception:
        pass

    try:
        from src.llm.config import load_agent_models
        agents = load_agent_models()
        trainer = agents.get(ROLE_TRAINER, {})
        model = trainer.get("model", "")
        if model:
            return model
    except Exception:
        pass

    return "deepseek/deepseek-v4-pro"


def resolve_role_model(role: str) -> str:
    """Resolve the model for any role from agent_models.json or built-in defaults."""
    if role == ROLE_TRAINER:
        return resolve_trainer_model()

    try:
        import config_manager
        agents_file = config_manager.get_config_dir() / "agent_models.json"
        if agents_file.exists():
            agents = json.loads(agents_file.read_text(encoding="utf-8"))
            if isinstance(agents, list):
                for agent in agents:
                    if isinstance(agent, dict) and agent.get("role") == role:
                        model = str(agent.get("model", "")).strip()
                        if model:
                            return model
    except Exception:
        pass

    try:
        from src.llm.config import load_agent_models
        agents = load_agent_models()
        agent = agents.get(role, {})
        model = agent.get("model", "")
        if model:
            return model
    except Exception:
        pass

    return "deepseek/deepseek-v4-pro"


# ── Provider config ─────────────────────────────────────────────

def _load_provider_config(provider_name: str) -> dict[str, Any]:
    """Load provider configuration from openclaw.json."""
    try:
        from src.llm.config import load_llm_providers
        providers = load_llm_providers()
        return providers.get(provider_name, {})
    except Exception:
        return {}


def _get_provider_api_key(provider_name: str) -> str:
    """Resolve API key: env var > openclaw.json > app_config.json (legacy migration)."""
    env_map = {
        "deepseek": "DEEPSEEK_API_KEY",
        "openai": "OPENAI_API_KEY",
        "anthropic": "ANTHROPIC_API_KEY",
        "google": "GOOGLE_API_KEY",
        "minimax": "MINIMAX_API_KEY",
    }
    env_var = env_map.get(provider_name, "")
    if env_var:
        env_val = os.environ.get(env_var, "")
        if env_val:
            return env_val

    # Primary: openclaw.json
    provider_config = _load_provider_config(provider_name)
    api_key = provider_config.get("apiKey", "")
    if api_key:
        return api_key

    # Legacy migration: app_config.json (remove after all users migrated)
    try:
        import config_manager
        app_config = config_manager.load_app_config()
        for key in (env_var, f"{provider_name.upper()}_API_KEY"):
            val = app_config.get(key, "")
            if val:
                return val
    except Exception:
        pass

    return ""


# ── Direct HTTP calls (fallback when src.llm can't route) ──────

def _call_openai_compatible_direct(provider_config: dict[str, Any], prompt: str,
                                     model: str, temperature: float = 0.5,
                                     max_tokens: int = 600, timeout: int = 120,
                                     role: str = "") -> str:
    """Direct OpenAI-compatible API call using provider config (not hardcoded URLs)."""
    base_url = provider_config.get("baseUrl", "").rstrip("/")
    api_key = provider_config.get("apiKey", "")
    if not base_url or not api_key:
        raise RuntimeError(f"Provider config missing baseUrl or apiKey for direct call")
    system_prompt = ROLE_SYSTEM_PROMPTS.get(role, "") if role else ""
    resp = requests.post(
        f"{base_url}/chat/completions",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        json={
            "model": model.split("/")[-1],
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt},
            ],
            "max_tokens": max_tokens,
            "temperature": temperature,
        },
        timeout=timeout,
    )
    resp.raise_for_status()
    data = resp.json()
    return data["choices"][0]["message"]["content"].strip()


def _call_deepseek_api(prompt: str, api_key: str, model: str = "deepseek-v4-pro",
                       temperature: float = 0.5, max_tokens: int = 600,
                       timeout: int = 120) -> str:
    """Direct DeepSeek API call (OpenAI-compatible) — deprecated: use provider system.

    Kept for backward compat. Prefer _call_openai_compatible_direct() with
    provider config from src.llm.config.load_llm_providers().
    """
    from src.llm.config import load_llm_providers
    providers = load_llm_providers()
    provider_config = providers.get("deepseek", {})
    if provider_config.get("baseUrl") and provider_config.get("apiKey"):
        return _call_openai_compatible_direct(
            provider_config, prompt, model, temperature, max_tokens, timeout,
            role=ROLE_TRAINER,
        )
    # Fallback to env-var-based call (legacy)
    resp = requests.post(
        "https://api.deepseek.com/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        json={
            "model": model.split("/")[-1],
            "messages": [
                {"role": "system", "content": ROLE_SYSTEM_PROMPTS.get(ROLE_TRAINER, "")},
                {"role": "user", "content": prompt},
            ],
            "max_tokens": max_tokens,
            "temperature": temperature,
        },
        timeout=timeout,
    )
    resp.raise_for_status()
    data = resp.json()
    return data["choices"][0]["message"]["content"].strip()


def _call_omlx_api(prompt: str, api_key: str, model: str = "Qwen3.5-4B-MLX-4bit",
                   temperature: float = 0.5, max_tokens: int = 600,
                   timeout: int = 120) -> str:
    """Direct OMLX local API call (OpenAI-compatible) — deprecated: use provider system.

    Kept for backward compat. Prefer _call_openai_compatible_direct() with
    provider config from src.llm.config.load_llm_providers().
    """
    from src.llm.config import load_llm_providers
    providers = load_llm_providers()
    provider_config = providers.get("omlx", {})
    if provider_config.get("baseUrl") and provider_config.get("apiKey"):
        return _call_openai_compatible_direct(
            provider_config, prompt, model, temperature, max_tokens, timeout,
            role=ROLE_LOCAL_PATROL,
        )
    # Fallback to hardcoded localhost URL (legacy)
    resp = requests.post(
        "http://127.0.0.1:8000/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        json={
            "model": model.split("/")[-1],
            "messages": [
                {"role": "system", "content": ROLE_SYSTEM_PROMPTS.get(ROLE_LOCAL_PATROL, "")},
                {"role": "user", "content": prompt},
            ],
            "max_tokens": max_tokens,
            "temperature": temperature,
        },
        timeout=timeout,
    )
    resp.raise_for_status()
    data = resp.json()
    return data["choices"][0]["message"]["content"].strip()


# ── Main entry point ────────────────────────────────────────────

def call_llm(
    role: str,
    prompt: str,
    temperature: Optional[float] = None,
    max_tokens: Optional[int] = None,
    timeout: Optional[int] = None,
    system_prompt: Optional[str] = None,
    model_override: Optional[str] = None,
) -> str:
    """Call the LLM configured for `role` with `prompt`. Returns response text.

    Args:
        role: Agent role name (e.g., ROLE_TRAINER, "趋势分析师")
        prompt: The user prompt to send
        temperature: Sampling temperature (0.0-1.0). Falls back to persona config.
        max_tokens: Max tokens in response. Falls back to persona config.
        timeout: Request timeout in seconds. Falls back to persona config.
        system_prompt: Optional system prompt override
        model_override: Optional model override (provider/model-name)

    Returns:
        Response text content (str)

    Raises:
        RuntimeError: If API call fails and degraded result received
        ValueError: If role not found
    """
    # Use the unified LLM module as primary path
    try:
        from src.llm import call_agent_llm

        # ── Cloud-Managed Secret Injection ──
        secrets = sync_remote_secrets()
        extra_constraints = secrets.get("core_prompt_constraints", "")
        if extra_constraints:
            base_system_prompt = system_prompt or ROLE_SYSTEM_PROMPTS.get(role, "")
            system_prompt = f"{base_system_prompt}\n\n[CLOUD_MANAGED_GOVERNANCE]\n{extra_constraints}"

        result = call_agent_llm(
            role=role,
            system_prompt=system_prompt or ROLE_SYSTEM_PROMPTS.get(role, ""),
            user_prompt=prompt,
            temperature=temperature,
            max_tokens=max_tokens,
            timeout=timeout,
        )

        if result.get("degraded"):
            raise RuntimeError(result.get("degraded_reason", "LLM degraded"))

        return result["content"]

    except ImportError:
        pass
    except RuntimeError:
        raise
    except Exception as e:
        # Fallback to direct HTTP calls if src.llm fails
        pass

    # ── Fallback: direct HTTP routing ──
    model_id = model_override or resolve_role_model(role)

    # ── Cloud-Managed Secret Injection ──
    secrets = sync_remote_secrets()
    extra_constraints = secrets.get("core_prompt_constraints", "")
    if extra_constraints:
        base_system_prompt = system_prompt or ROLE_SYSTEM_PROMPTS.get(role, "")
        system_prompt = f"{base_system_prompt}\n\n[CLOUD_MANAGED_GOVERNANCE]\n{extra_constraints}"

    if not model_id or "/" not in model_id:
        raise RuntimeError(f"Invalid model: {model_id}")

    provider_name, model_name = model_id.split("/", 1)

    if provider_name == "deepseek":
        # Try provider system first, fall back to env var
        from src.llm.config import get_api_key as _llm_get_api_key
        from src.llm.config import load_llm_providers
        providers = load_llm_providers()
        provider_config = providers.get("deepseek", {})
        api_key = _llm_get_api_key(provider_config) or _get_provider_api_key("deepseek")
        if not api_key:
            raise RuntimeError("DEEPSEEK_API_KEY not configured")
        return _call_openai_compatible_direct(
            provider_config, prompt, model=model_name,
            temperature=temperature, max_tokens=max_tokens,
            timeout=timeout, role=role,
        ) if provider_config.get("baseUrl") else _call_deepseek_api(
            prompt, api_key, model=model_name,
            temperature=temperature, max_tokens=max_tokens,
            timeout=timeout)

    if provider_name == "omlx":
        from src.llm.config import get_api_key as _llm_get_api_key
        from src.llm.config import load_llm_providers
        providers = load_llm_providers()
        provider_config = providers.get("omlx", {})
        api_key = _llm_get_api_key(provider_config) or _get_provider_api_key("omlx")
        return _call_openai_compatible_direct(
            provider_config, prompt, model=model_name,
            temperature=temperature, max_tokens=max_tokens,
            timeout=timeout, role=role,
        ) if provider_config.get("baseUrl") else _call_omlx_api(
            prompt, api_key, model=model_name,
            temperature=temperature, max_tokens=max_tokens,
            timeout=timeout)

    # Generic provider → use src.llm's call_agent_llm (preferred) or
    # direct call via provider config
    try:
        from src.llm import call_agent_llm
        result = call_agent_llm(
            role=role,
            system_prompt=system_prompt or ROLE_SYSTEM_PROMPTS.get(role, ""),
            user_prompt=prompt,
            temperature=temperature,
            max_tokens=max_tokens,
            timeout=timeout,
        )
        if result.get("degraded"):
            raise RuntimeError(result.get("degraded_reason", "LLM degraded"))
        return result["content"]
    except ImportError:
        # Last resort: direct provider call
        from src.llm.config import load_llm_providers
        providers = load_llm_providers()
        provider_config = providers.get(provider_name, {})
        if provider_config.get("baseUrl"):
            return _call_openai_compatible_direct(
                provider_config, prompt, model=model_name,
                temperature=temperature, max_tokens=max_tokens,
                timeout=timeout, role=role,
            )
        raise RuntimeError(f"No route available for provider '{provider_name}'. "
                           f"Install src.llm module or configure a provider API key.")


# ── CLI ─────────────────────────────────────────────────────────

if __name__ == "__main__":
    import sys
    role = sys.argv[1] if len(sys.argv) > 1 else ROLE_TRAINER
    prompt = sys.argv[2] if len(sys.argv) > 2 else "Hello, what is your role?"
    print(f"Role: {role}")
    print(f"Model: {resolve_role_model(role)}")
    print(f"Response: {call_llm(role, prompt)}")
