"""AI 公共调用工具 — 支持 deepseek / minimax / ollama / openai_compat"""
import re
import json
from typing import AsyncGenerator
import httpx
from fastapi import HTTPException
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

AI_PROVIDER_DEFAULTS: dict[str, dict] = {
    "deepseek":      {"base_url": "https://api.deepseek.com",  "model": "deepseek-chat"},
    "minimax":       {"base_url": "https://api.minimaxi.com",  "model": "MiniMax-M2.7"},
    "ollama":        {"base_url": "http://localhost:11434",     "model": "llama3"},
    "openai_compat": {"base_url": "https://api.openai.com",    "model": "gpt-3.5-turbo"},
}


async def resolve_ai_extra(db: AsyncSession, tenant_id: int, tenant_extra: dict) -> dict:
    """
    按优先级解析出实际调用 AI 时使用的配置：
    1. 租户被允许自定义（ai_allow_custom=True）且已自行配置好（ai_api_key，或 provider=ollama 且 base_url/model 至少填了一个）→ 用租户自己的配置
    2. 否则用 superadmin 为该租户分配的 AiProfile（tenants.ai_profile_id）
    3. 否则回退到标记为 is_default=True 且 is_active=True 的 AiProfile
    4. 都没有 → 原样返回 tenant_extra（call_ai 会走现有的"未配置"报错）
    """
    from app.core.models.tenant import Tenant
    from app.core.models.ai_profile import AiProfile

    tenant = (await db.execute(
        select(Tenant.ai_allow_custom, Tenant.ai_profile_id).where(Tenant.id == tenant_id)
    )).first()
    if tenant is None:
        return tenant_extra

    allow_custom, profile_id = tenant
    is_ollama_configured = tenant_extra.get("ai_provider") == "ollama" and bool(
        tenant_extra.get("ai_base_url") or tenant_extra.get("ai_model")
    )
    has_custom_key = bool(tenant_extra.get("ai_api_key")) or is_ollama_configured
    if allow_custom and has_custom_key:
        return tenant_extra

    profile = None
    if profile_id:
        profile = (await db.execute(
            select(AiProfile).where(AiProfile.id == profile_id, AiProfile.is_active == True)
        )).scalar_one_or_none()
    if profile is None:
        profile = (await db.execute(
            select(AiProfile).where(AiProfile.is_default == True, AiProfile.is_active == True)
        )).scalars().first()

    if profile is None:
        return tenant_extra
    return {**tenant_extra, **profile.to_ai_extra()}


async def call_ai(
    user_prompt: str,
    extra: dict,
    *,
    system_prompt: str = "",
    max_tokens: int = 2048,
    timeout: int = 120,
) -> str:
    """
    使用租户 AI 配置（extra 字段）调用 AI，返回纯文本回复。
    extra 应包含：ai_provider, ai_api_key, ai_base_url, ai_model
    """
    provider = extra.get("ai_provider", "deepseek")
    api_key  = extra.get("ai_api_key", "")
    base_url = extra.get("ai_base_url") or AI_PROVIDER_DEFAULTS.get(provider, {}).get("base_url", "")
    model    = extra.get("ai_model")    or AI_PROVIDER_DEFAULTS.get(provider, {}).get("model", "")

    if not base_url:
        raise HTTPException(status_code=400, detail="AI 服务地址未配置，请在系统设置 → AI 配置中完善")
    if provider != "ollama" and not api_key:
        raise HTTPException(status_code=400, detail="API Key 未配置，请在系统设置 → AI 配置中完善")

    messages: list[dict] = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": user_prompt})

    # 去掉用户可能多填的 /v1 /v2 等路径后缀，统一由代码拼接
    base_clean = re.sub(r'/v\d+/?$', '', base_url.rstrip('/'))

    if provider == "ollama":
        url     = base_clean + "/api/chat"
        headers = {"Content-Type": "application/json"}
        payload: dict = {"model": model, "messages": messages, "stream": False}
    else:
        url     = base_clean + "/v1/chat/completions"
        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
        payload = {"model": model, "messages": messages, "max_tokens": max_tokens}

    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.post(url, json=payload, headers=headers)
            resp.raise_for_status()
            data = resp.json()
    except httpx.TimeoutException:
        raise HTTPException(status_code=504, detail="AI 响应超时，请稍后重试或换更小的模型")
    except httpx.HTTPStatusError as e:
        raise HTTPException(status_code=502, detail=f"AI 服务错误 {e.response.status_code}：{e.response.text[:200]}")
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"连接 AI 失败：{str(e)}")

    try:
        text = data["message"]["content"] if provider == "ollama" else data["choices"][0]["message"]["content"]
    except (KeyError, IndexError):
        text = str(data)
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()


async def call_ai_stream(
    messages: list[dict],
    extra: dict,
    *,
    max_tokens: int = 1024,
    timeout: int = 60,
) -> AsyncGenerator[str, None]:
    """流式调用 AI，逐 token yield delta 字符串。messages 已包含 system message。"""
    provider = extra.get("ai_provider", "deepseek")
    api_key  = extra.get("ai_api_key", "")
    base_url = extra.get("ai_base_url") or AI_PROVIDER_DEFAULTS.get(provider, {}).get("base_url", "")
    model    = extra.get("ai_model")    or AI_PROVIDER_DEFAULTS.get(provider, {}).get("model", "")

    if not base_url:
        yield "[ERROR] AI 服务地址未配置，请在系统设置 → AI 配置中完善"
        return
    if provider != "ollama" and not api_key:
        yield "[ERROR] API Key 未配置，请在系统设置 → AI 配置中完善"
        return

    base_clean = re.sub(r'/v\d+/?$', '', base_url.rstrip('/'))

    if provider == "ollama":
        url     = base_clean + "/api/chat"
        headers = {"Content-Type": "application/json"}
        payload: dict = {"model": model, "messages": messages, "stream": True}
    else:
        url     = base_clean + "/v1/chat/completions"
        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
        payload = {"model": model, "messages": messages, "max_tokens": max_tokens, "stream": True}

    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            async with client.stream("POST", url, json=payload, headers=headers) as resp:
                resp.raise_for_status()
                async for line in resp.aiter_lines():
                    if not line.startswith("data: "):
                        continue
                    data_str = line[6:].strip()
                    if data_str == "[DONE]":
                        return
                    try:
                        data = json.loads(data_str)
                        if provider == "ollama":
                            delta = data.get("message", {}).get("content", "") or ""
                            if data.get("done"):
                                return
                        else:
                            delta = data["choices"][0]["delta"].get("content", "") or ""
                        if delta:
                            yield delta
                    except (json.JSONDecodeError, KeyError, IndexError):
                        continue
    except httpx.TimeoutException:
        yield "[ERROR] AI 响应超时，请稍后重试"
    except Exception as e:
        yield f"[ERROR] 连接 AI 失败：{str(e)[:100]}"


def _partial_tag_suffix(s: str, tag: str) -> int:
    """Return length of longest suffix of s that is a prefix of tag."""
    for i in range(min(len(tag) - 1, len(s)), 0, -1):
        if s.endswith(tag[:i]):
            return i
    return 0


async def strip_think_blocks(stream: AsyncGenerator[str, None]) -> AsyncGenerator[str, None]:
    """Filter <think>...</think> blocks from a streaming token generator."""
    in_think = False
    buf = ""

    async for chunk in stream:
        if not chunk:
            continue
        buf += chunk
        out = ""

        while buf:
            if in_think:
                idx = buf.find("</think>")
                if idx == -1:
                    hold = _partial_tag_suffix(buf, "</think>")
                    buf = buf[-hold:] if hold else ""
                    break
                else:
                    buf = buf[idx + 8:]
                    in_think = False
            else:
                idx = buf.find("<think>")
                if idx == -1:
                    hold = _partial_tag_suffix(buf, "<think>")
                    out += buf[:-hold] if hold else buf
                    buf = buf[-hold:] if hold else ""
                    break
                else:
                    out += buf[:idx]
                    buf = buf[idx + 7:]
                    in_think = True

        if out:
            yield out

    if buf and not in_think:
        yield buf