"""
LLM bridge for pipeline handlers.

Provides a simple async interface for handlers to call LLM APIs.
Looks up model config from the llm database table first,
falls back to environment variables.
"""

import json
import logging
import os

logger = logging.getLogger("pipeline.llm_bridge")

# Cache: model_name -> {api_base, api_key, model_id}
_model_cache: dict = {}


def _decrypt_key(encrypted: str) -> str:
    """Decrypt api_key stored with password_encode. Falls back to plaintext."""
    if not encrypted:
        return ""
    try:
        from appPublic.rc4 import unpassword
        from appPublic.jsonConfig import getConfig
        config = getConfig()
        # 与 ahserver.globalEnv.get_password_key() 一致：password_key 为空时用默认 key
        key = config.password_key or 'QRIVSRHrthhwyjy176556332'
        # unpassword(code, key)：code=密文，key=密钥（之前参数顺序写反了）
        return unpassword(encrypted, key)
    except Exception:
        return encrypted  # already plaintext or decrypt failed


async def _get_model_config(model_name: str = None, org_id: str = None) -> dict:
    """Look up model config from llm table. Returns dict with api_base, api_key, model_id.

    org_id 多租户隔离：非空时只取本机构(org_id=${org}$)的模型，不含系统级兜底；
    org_id 为空时不过滤（向后兼容）。
    """
    global _model_cache

    cache_key = f"{model_name or ''}:{org_id or ''}"
    if model_name and cache_key in _model_cache:
        return _model_cache[cache_key]

    try:
        from sqlor.dbpools import DBPools
        db = DBPools()
        dbname = "pipeline"
        async with db.sqlorContext(dbname) as sor:
            if model_name:
                sql = "SELECT api_base, api_key, model_id FROM llm WHERE name=${name}$ AND status='active'"
                params = {"name": model_name}
                if org_id:
                    sql += " AND org_id=${org}$"
                    params["org"] = org_id
                sql += " LIMIT 1"
                recs = await sor.sqlExe(sql, params)
            else:
                sql = "SELECT api_base, api_key, model_id, name FROM llm WHERE status='active'"
                params = {}
                if org_id:
                    sql += " AND org_id=${org}$"
                    params["org"] = org_id
                sql += " ORDER BY id LIMIT 1"
                recs = await sor.sqlExe(sql, params)
            if recs:
                r = recs[0]
                cfg = {
                    "api_base": getattr(r, "api_base", "") or "",
                    "api_key": _decrypt_key(getattr(r, "api_key", "") or ""),
                    "model_id": getattr(r, "model_id", "") or "",
                }
                _model_cache[cache_key] = cfg
                return cfg
    except Exception as e:
        logger.warning("llm_bridge: DB lookup failed: %s", e)

    return {}


# 瞬时错误重试：超时/连接错误/限流/服务端5xx 均重试；配置错误(无key)、鉴权/参数4xx 不重试
_RETRYABLE_STATUS = (429, 500, 502, 503, 504)
_LLM_MAX_ATTEMPTS = 3
# 大上下文（角色 agent 多轮累积）生成偏慢，180s 曾触发超时（空 err=），放宽到 5 分钟；
# connect 单独设 30s，连接建立失败能快速失败并重试，而不是干等 5 分钟。
_LLM_TOTAL_TIMEOUT = 300
_LLM_CONNECT_TIMEOUT = 30


async def _post_chat_completion(url: str, headers: dict, payload: dict) -> dict:
    """POST /chat/completions，带瞬时错误重试。返回解析后的 JSON dict。"""
    import aiohttp
    import asyncio

    last_exc = None
    for attempt in range(_LLM_MAX_ATTEMPTS):
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    url, headers=headers, json=payload,
                    timeout=aiohttp.ClientTimeout(total=_LLM_TOTAL_TIMEOUT, connect=_LLM_CONNECT_TIMEOUT),
                ) as resp:
                    if resp.status != 200:
                        text = await resp.text()
                        err = ValueError("LLM API error %d: %s" % (resp.status, text[:300]))
                        if resp.status in _RETRYABLE_STATUS and attempt < _LLM_MAX_ATTEMPTS - 1:
                            last_exc = err
                            await asyncio.sleep(2 * (attempt + 1))
                            continue
                        raise err
                    data = await resp.json()
                    # 网关偶发返回 200 但内容无 choices（错误 JSON），视为可重试的瞬时异常；
                    # 不处理会导致上层 llm_call_msgs_native 的 data["choices"] 抛 KeyError。
                    if not isinstance(data, dict) or "choices" not in data:
                        err = ValueError("LLM 响应缺 choices: %s" % json.dumps(data, ensure_ascii=False)[:300])
                        if attempt < _LLM_MAX_ATTEMPTS - 1:
                            last_exc = err
                            await asyncio.sleep(2 * (attempt + 1))
                            continue
                        raise err
                    return data
        except (asyncio.TimeoutError, aiohttp.ClientError) as e:
            last_exc = e
            if attempt < _LLM_MAX_ATTEMPTS - 1:
                logger.warning("llm_bridge: 瞬时错误重试 %d/%d: %s", attempt + 1, _LLM_MAX_ATTEMPTS, e)
                await asyncio.sleep(2 * (attempt + 1))
                continue
            raise

    raise last_exc if last_exc else ValueError("LLM call failed")


async def llm_call(prompt: str, model: str = None, temperature: float = 0.7, org_id: str = None) -> str:
    """Call LLM and return text response.

    Backend priority:
    1. harnessed_agent.llm_chat (if loaded in ServerEnv)
    2. DB llm table (api_base + api_key)
    3. Environment variables (LLM_API_BASE, LLM_API_KEY, LLM_MODEL)

    org_id 多租户隔离：非空时 llm 表查询只取「本机构 + 系统级」模型。
    """
    # Priority 1: harnessed_agent
    try:
        from ahserver.serverenv import ServerEnv
        env = ServerEnv()
        if hasattr(env, 'llm_chat'):
            result = await env.llm_chat(prompt, model=model, temperature=temperature)
            if isinstance(result, dict):
                return result.get("content", result.get("text", str(result)))
            return str(result)
    except Exception:
        pass

    # Priority 2: DB llm table
    cfg = await _get_model_config(model, org_id=org_id)
    if cfg.get("api_key") and cfg.get("api_base"):
        api_base = cfg["api_base"]
        api_key = cfg["api_key"]
        model_id = cfg.get("model_id") or model or "default"
        logger.info("llm_bridge: using DB model config for %s -> %s", model, api_base)
    else:
        # Priority 3: Environment variables
        api_base = os.environ.get("LLM_API_BASE", "https://api.openai.com/v1")
        api_key = os.environ.get("LLM_API_KEY", "")
        model_id = model or os.environ.get("LLM_MODEL", "gpt-4o-mini")

    if not api_key:
        raise ValueError("No LLM API configured. Please add a model in the llm table or set LLM_API_KEY env var.")

    import aiohttp

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": model_id,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": temperature,
    }

    url = api_base.rstrip("/") + "/chat/completions"
    data = await _post_chat_completion(url, headers, payload)
    return data["choices"][0]["message"]["content"]


async def call_llm(tenant_id: str, prompt: str, model: str = None, temperature: float = 0.7) -> str:
    """SDLC handler interface — delegates to llm_call."""
    return await llm_call(prompt, model=model, temperature=temperature)


async def llm_call_msgs(messages: list, model: str = None, temperature: float = 0.7, org_id: str = None) -> str:
    """Call LLM with full message array (system/user/assistant)."""
    import aiohttp

    cfg = await _get_model_config(model, org_id=org_id)
    if cfg.get("api_key") and cfg.get("api_base"):
        api_base = cfg["api_base"]
        api_key = cfg["api_key"]
        model_id = cfg.get("model_id") or model or "default"
    else:
        api_base = os.environ.get("LLM_API_BASE", "https://api.openai.com/v1")
        api_key = os.environ.get("LLM_API_KEY", "")
        model_id = model or os.environ.get("LLM_MODEL", "gpt-4o-mini")

    if not api_key:
        raise ValueError("No LLM API configured")

    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    payload = {"model": model_id, "messages": messages, "temperature": temperature}

    url = api_base.rstrip("/") + "/chat/completions"
    data = await _post_chat_completion(url, headers, payload)
    return data["choices"][0]["message"]["content"]


async def llm_call_msgs_native(messages: list, tools: list = None, model: str = None, temperature: float = 0.7, org_id: str = None) -> dict:
    """Native function calling. 传入 tools JSON schema，返回 message dict。

    Returns:
        {"content": str, "tool_calls": [{"id","type","function":{"name","arguments"}}]}
    当模型返回 tool_calls 时，content 通常为空字符串。
    """
    import aiohttp

    cfg = await _get_model_config(model, org_id=org_id)
    if cfg.get("api_key") and cfg.get("api_base"):
        api_base = cfg["api_base"]
        api_key = cfg["api_key"]
        model_id = cfg.get("model_id") or model or "default"
    else:
        api_base = os.environ.get("LLM_API_BASE", "https://api.openai.com/v1")
        api_key = os.environ.get("LLM_API_KEY", "")
        model_id = model or os.environ.get("LLM_MODEL", "gpt-4o-mini")

    if not api_key:
        raise ValueError("No LLM API configured")

    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    payload = {"model": model_id, "messages": messages, "temperature": temperature}
    if tools:
        payload["tools"] = tools
        payload["tool_choice"] = "auto"

    url = api_base.rstrip("/") + "/chat/completions"
    data = await _post_chat_completion(url, headers, payload)
    msg = data["choices"][0]["message"]
    return {
        "content": msg.get("content") or "",
        "tool_calls": msg.get("tool_calls") or [],
    }