"""
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 = {}


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

    if model_name and model_name in _model_cache:
        return _model_cache[model_name]

    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' LIMIT 1"
                recs = await sor.sqlExe(sql, {"name": model_name})
            else:
                sql = "SELECT api_base, api_key, model_id, name FROM llm WHERE status='active' ORDER BY id LIMIT 1"
                recs = await sor.sqlExe(sql, {})
            if recs:
                r = recs[0]
                cfg = {
                    "api_base": getattr(r, "api_base", "") or "",
                    "api_key": getattr(r, "api_key", "") or "",
                    "model_id": getattr(r, "model_id", "") or "",
                }
                cache_key = model_name or getattr(r, "name", "")
                if cache_key:
                    _model_cache[cache_key] = cfg
                return cfg
    except Exception as e:
        logger.warning("llm_bridge: DB lookup failed: %s", e)

    return {}


async def llm_call(prompt: str, model: str = None, temperature: float = 0.7) -> 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)
    """
    # 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)
    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"
    async with aiohttp.ClientSession() as session:
        async with session.post(
            url, headers=headers, json=payload,
            timeout=aiohttp.ClientTimeout(total=120)
        ) as resp:
            if resp.status != 200:
                text = await resp.text()
                raise ValueError(f"LLM API error {resp.status}: {text[:300]}")
            data = await resp.json()
            return data["choices"][0]["message"]["content"]