"""内部 LLM 代理：短期 token 鉴权，真 api_key 永不出服务进程。

背景（设计铁律）：
    运行环境（角色 agent 的 run_shell 子进程、bwrap 沙箱、远程 SSH 目标机、被开发的应用）
    里的东西需要调 LLM 时，**绝不下发真 api_key**——运行环境恰恰是 agent 能自由执行
    shell 的地方，key 落进去等于可被 `env | grep KEY` 读出、写进交付件、commit 进 git、
    打进日志；落磁盘（env/*.json 会 git 提交）更会泄露给所有能读该仓库的人，并破坏
    llm 表的机构隔离。

做法：
    1. 服务端签发短期 token（绑定 org_id + project_id + task_id，有过期时间/调用上限，可吊销）。
    2. 运行环境拿 token 当 api_key，base_url 指向本平台的 OpenAI 兼容代理端点。
    3. 代理端点校验 token → 按 token 绑定的 org_id 走 llm_bridge 解析真 key（机构隔离）
       → 转发上游 → 回填用量。真 key 只在服务进程内存里出现。

适配 browser-use / LangChain / OpenAI SDK 等一切 OpenAI 兼容客户端：
    ChatOpenAI(base_url="http://<host>/pipeline_core/api/llm_proxy_v1/chat/completions" 的父路径,
               api_key="<短期 token>")
"""

import json
import logging
import secrets
from datetime import datetime, timedelta

from sqlor.dbpools import DBPools
from appPublic.uniqueID import getID

DBNAME = "pipeline"
logger = logging.getLogger("pipeline.llm_proxy")

TOKEN_PREFIX = "plk-"       # pipeline llm key
S_ACTIVE = "active"
S_REVOKED = "revoked"

DEFAULT_TTL_HOURS = 8
DEFAULT_MAX_CALLS = 500

# 失败限速：防 token 枚举。按来源标识（token 前缀/调用方）计数，
# 窗口内失败次数超阈值则直接拒绝，不再查库（避免枚举放大 DB 压力）。
_FAIL_WINDOW_SEC = 300
_FAIL_MAX = 20
_fail_counters: dict = {}


def _fail_key(token):
    """限速键：取 token 前 12 字符（足够区分调用方，又不在内存留全量密钥）。"""
    t = (token or '').strip()
    if t.lower().startswith('bearer '):
        t = t[7:].strip()
    return t[:12] or 'anon'


def _rate_limited(token):
    """是否已被限速。顺带清理过期窗口。"""
    import time
    now = time.time()
    k = _fail_key(token)
    rec = _fail_counters.get(k)
    if not rec:
        return False
    start, cnt = rec
    if now - start > _FAIL_WINDOW_SEC:
        _fail_counters.pop(k, None)
        return False
    return cnt >= _FAIL_MAX


def _record_fail(token):
    """记一次鉴权失败。"""
    import time
    now = time.time()
    k = _fail_key(token)
    rec = _fail_counters.get(k)
    if not rec or now - rec[0] > _FAIL_WINDOW_SEC:
        _fail_counters[k] = (now, 1)
    else:
        _fail_counters[k] = (rec[0], rec[1] + 1)
    # 防内存膨胀：条目过多时清理已过期窗口
    if len(_fail_counters) > 5000:
        for kk in [x for x, v in _fail_counters.items() if now - v[0] > _FAIL_WINDOW_SEC]:
            _fail_counters.pop(kk, None)


def _get_db():
    db = DBPools()
    if not db.databases:
        from appPublic.jsonConfig import getConfig
        config = getConfig()
        if config.databases:
            db.databases = config.databases
    return db, DBNAME


def _now():
    return datetime.now()


async def create_llm_token(org_id, project_id="", task_id="", model_name="",
                           purpose="", ttl_hours=DEFAULT_TTL_HOURS,
                           max_calls=DEFAULT_MAX_CALLS, created_by=""):
    """签发短期 LLM 代理 token。返回 (True, token) 或 (False, 错误信息)。

    org_id 决定机构隔离边界：代理调用时按此 org_id 查 llm 表解析真 key，
    运行环境无法越权使用其他机构的模型。
    """
    if not org_id:
        return False, "缺少 org_id（机构隔离必需）"
    try:
        ttl = int(ttl_hours)
    except (TypeError, ValueError):
        ttl = DEFAULT_TTL_HOURS
    if ttl <= 0 or ttl > 72:
        ttl = DEFAULT_TTL_HOURS
    try:
        mc = int(max_calls)
    except (TypeError, ValueError):
        mc = DEFAULT_MAX_CALLS
    if mc <= 0:
        mc = DEFAULT_MAX_CALLS

    token = TOKEN_PREFIX + secrets.token_urlsafe(32)
    expires_at = _now() + timedelta(hours=ttl)
    db, dbname = _get_db()
    async with db.sqlorContext(dbname) as sor:
        await sor.C('pipeline_llm_tokens', {
            'id': getID(),
            'token': token,
            'org_id': org_id,
            'project_id': project_id or '',
            'task_id': task_id or '',
            'model_name': model_name or '',
            'purpose': purpose or '',
            'status': S_ACTIVE,
            'expires_at': expires_at.strftime('%Y-%m-%d %H:%M:%S'),
            'max_calls': mc,
            'call_count': 0,
            'prompt_tokens': 0,
            'completion_tokens': 0,
            'created_by': created_by or '',
        })
        await sor.sqlExe("COMMIT", {})
    logger.info("create_llm_token: org=%s project=%s task=%s purpose=%s ttl=%dh max_calls=%d",
                org_id, project_id, task_id, purpose, ttl, mc)
    return True, token


async def verify_llm_token(token):
    """校验 token。返回 (True, 绑定信息 dict) 或 (False, 错误信息)。

    校验项：存在、status=active、未过期、未超调用上限。
    """
    if not token:
        return False, "缺少 token"
    token = token.strip()
    if token.lower().startswith("bearer "):
        token = token[7:].strip()
    db, dbname = _get_db()
    async with db.sqlorContext(dbname) as sor:
        recs = await sor.sqlExe(
            "SELECT id, org_id, project_id, task_id, model_name, status, expires_at, "
            "max_calls, call_count FROM pipeline_llm_tokens WHERE token=${t}$",
            {"t": token})
        await sor.sqlExe("COMMIT", {})
        if not recs:
            return False, "token 无效"
        r = recs[0]
        status = getattr(r, 'status', '') or ''
        if status != S_ACTIVE:
            return False, f"token 已{('吊销' if status == S_REVOKED else status)}"
        exp = getattr(r, 'expires_at', None)
        if exp:
            if isinstance(exp, str):
                try:
                    exp = datetime.strptime(exp[:19], '%Y-%m-%d %H:%M:%S')
                except ValueError:
                    exp = None
            if exp and _now() > exp:
                return False, "token 已过期"
        mc = int(getattr(r, 'max_calls', 0) or 0)
        cc = int(getattr(r, 'call_count', 0) or 0)
        if mc and cc >= mc:
            return False, f"token 调用次数已达上限（{mc}）"
        return True, {
            "id": getattr(r, 'id', ''),
            "org_id": getattr(r, 'org_id', '') or '',
            "project_id": getattr(r, 'project_id', '') or '',
            "task_id": getattr(r, 'task_id', '') or '',
            "model_name": getattr(r, 'model_name', '') or '',
        }


async def revoke_llm_token(token=None, token_id=None):
    """吊销 token（按 token 明文或 id）。返回 (True, 消息) 或 (False, 错误)。"""
    if not token and not token_id:
        return False, "须指定 token 或 token_id"
    db, dbname = _get_db()
    async with db.sqlorContext(dbname) as sor:
        if token_id:
            await sor.sqlExe(
                "UPDATE pipeline_llm_tokens SET status=${s}$ WHERE id=${i}$",
                {"s": S_REVOKED, "i": token_id})
        else:
            t = token.strip()
            if t.lower().startswith("bearer "):
                t = t[7:].strip()
            await sor.sqlExe(
                "UPDATE pipeline_llm_tokens SET status=${s}$ WHERE token=${t}$",
                {"s": S_REVOKED, "t": t})
        await sor.sqlExe("COMMIT", {})
    return True, "已吊销"


async def revoke_task_tokens(task_id):
    """任务结束时吊销该任务签发的全部 token（防止 token 生命周期超出任务）。"""
    if not task_id:
        return False, "缺少 task_id"
    db, dbname = _get_db()
    async with db.sqlorContext(dbname) as sor:
        await sor.sqlExe(
            "UPDATE pipeline_llm_tokens SET status=${s}$ "
            "WHERE task_id=${t}$ AND status=${a}$",
            {"s": S_REVOKED, "t": task_id, "a": S_ACTIVE})
        await sor.sqlExe("COMMIT", {})
    return True, "已吊销该任务全部 token"


async def _record_usage(token_id, usage):
    """回填用量（调用次数 + token 数）。usage 为上游返回的 usage dict。"""
    if not token_id:
        return
    pt = 0
    ct = 0
    if isinstance(usage, dict):
        try:
            pt = int(usage.get('prompt_tokens') or 0)
            ct = int(usage.get('completion_tokens') or 0)
        except (TypeError, ValueError):
            pt, ct = 0, 0
    db, dbname = _get_db()
    try:
        async with db.sqlorContext(dbname) as sor:
            await sor.sqlExe(
                "UPDATE pipeline_llm_tokens SET call_count=call_count+1, "
                "prompt_tokens=prompt_tokens+${pt}$, completion_tokens=completion_tokens+${ct}$, "
                "last_used_at=NOW() WHERE id=${i}$",
                {"pt": pt, "ct": ct, "i": token_id})
            await sor.sqlExe("COMMIT", {})
    except Exception as e:
        logger.warning("_record_usage failed: %s", e)


async def proxy_chat_completion(token, payload):
    """OpenAI 兼容代理转发。

    token  运行环境持有的短期 token（当 api_key 用）
    payload 客户端原始请求体（含 model/messages/tools/temperature 等）

    返回 (True, 上游响应 dict) 或 (False, 错误信息)。
    真 api_key 由 llm_bridge 按 token 绑定的 org_id 从 llm 表解析，不下发给调用方。
    """
    # 失败限速前置：窗口内鉴权失败过多直接拒绝（防 token 枚举）
    if _rate_limited(token):
        logger.warning("proxy_chat_completion: 限速拒绝 key=%s", _fail_key(token))
        return False, "鉴权失败次数过多，请稍后再试"
    ok, info = await verify_llm_token(token)
    if not ok:
        _record_fail(token)
        return False, info
    if not isinstance(payload, dict) or not payload.get('messages'):
        return False, "请求体缺 messages"

    org_id = info['org_id']
    # 模型选择：token 绑定了 model_name 则强制用它（防运行环境越权指定贵模型）；
    # 否则用请求里的 model；都没有则由 llm_bridge 取该机构第一个 active 模型。
    model_name = info.get('model_name') or payload.get('model') or None

    from .llm_bridge import _get_model_config, _post_chat_completion
    cfg = await _get_model_config(model_name, org_id=org_id)
    if not (cfg.get('api_key') and cfg.get('api_base')):
        return False, f"机构 {org_id} 未配置可用模型（llm 表 status=active）"

    api_base = cfg['api_base']
    api_key = cfg['api_key']          # 只在本进程内存中使用，不返回给调用方
    model_id = cfg.get('model_id') or model_name or 'default'

    # 透传客户端参数（messages/tools/temperature/max_tokens 等），但 model 换成真实 model_id，
    # 且不透传 stream（代理暂不支持流式）。
    upstream = {k: v for k, v in payload.items() if k not in ('model', 'stream')}
    upstream['model'] = model_id

    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    url = api_base.rstrip('/') + '/chat/completions'
    try:
        data = await _post_chat_completion(url, headers, upstream)
    except Exception as e:
        logger.warning("proxy_chat_completion upstream failed: %s", e)
        return False, f"上游调用失败: {str(e)[:300]}"

    await _record_usage(info['id'], data.get('usage'))
    logger.info("proxy_chat_completion: org=%s project=%s model=%s ok",
                org_id, info.get('project_id'), model_id)
    return True, data


async def list_llm_tokens(org_id=None, project_id=None, status=None, limit=50):
    """列出已签发 token（不返回 token 明文，只返回元信息 + 用量，供审计/管理）。"""
    db, dbname = _get_db()
    async with db.sqlorContext(dbname) as sor:
        conds = ["1=1"]
        params = {}
        if org_id:
            conds.append("org_id=${o}$")
            params["o"] = org_id
        if project_id:
            conds.append("project_id=${p}$")
            params["p"] = project_id
        if status:
            conds.append("status=${s}$")
            params["s"] = status
        where = " AND ".join(conds)
        recs = await sor.sqlExe(
            f"SELECT id, org_id, project_id, task_id, model_name, purpose, status, "
            f"expires_at, max_calls, call_count, prompt_tokens, completion_tokens, "
            f"created_by, created_at, last_used_at FROM pipeline_llm_tokens "
            f"WHERE {where} ORDER BY created_at DESC LIMIT {int(limit)}", params)
        await sor.sqlExe("COMMIT", {})
        out = []
        for r in (recs or []):
            try:
                out.append(dict(r))
            except (TypeError, ValueError):
                pass
        return out


def load_llm_proxy():
    """注册到 ServerEnv，供 dspy 端点调用。"""
    from ahserver.serverenv import ServerEnv
    env = ServerEnv()
    env.create_llm_token = create_llm_token
    env.verify_llm_token = verify_llm_token
    env.revoke_llm_token = revoke_llm_token
    env.revoke_task_tokens = revoke_task_tokens
    env.proxy_chat_completion = proxy_chat_completion
    env.list_llm_tokens = list_llm_tokens