from __future__ import annotations

import json
import re
from typing import Any

import httpx

from app.config import settings
from app.models import JointDef, LinkDef, RobotDraft


SYSTEM_PROMPT = """你是机器人 URDF 运动学专家。根据零件列表与包围盒，直接设计 link/joint 树。

约定：
- 单位：origin_xyz 用毫米；origin_rpy 用弧度。
- 第一个关节相对世界系，origin 靠近对应零件 center，不要写成 [0,0,0]。
- part_names 只能用输入里的名字（如 Solid_0），禁止发明新名字。
- 每个非 base_link 恰好一个 parent joint。
- axis 为单位向量；不要把所有轴都写成 [0,0,1]。
- rationale 最多 20 字。
- 只输出一个合法 JSON 对象，不要 Markdown、不要注释、不要尾逗号。

JSON schema:
{"robot_name":"string","links":[{"name":"string","part_names":["string"]}],"joints":[{"name":"string","joint_type":"revolute|prismatic|fixed","parent":"string","child":"string","origin_xyz":[0,0,0],"origin_rpy":[0,0,0],"axis":[0,0,1],"lower":-3.14,"upper":3.14,"effort":100,"velocity":1,"rationale":"string"}],"notes":["string"]}
"""


def _strip_fences(text: str) -> str:
    text = text.strip()
    if text.startswith("```"):
        text = re.sub(r"^```(?:json)?\s*", "", text)
        text = re.sub(r"\s*```$", "", text)
    return text.strip()


def _repair_json_text(text: str) -> str:
    """Best-effort cleanup for common LLM JSON issues / truncation."""
    text = _strip_fences(text)
    # Chinese quotes → ASCII
    text = text.replace("“", '"').replace("”", '"').replace("‘", "'").replace("’", "'")
    # Remove // comments
    text = re.sub(r"(?m)^\s*//.*?$", "", text)
    # Trailing commas before } or ]
    text = re.sub(r",\s*([}\]])", r"\1", text)

    # If truncated, close open braces/brackets
    start = text.find("{")
    if start < 0:
        return text
    text = text[start:]
    # Drop incomplete trailing string if odd number of unescaped quotes in last line
    in_str = False
    escape = False
    stack: list[str] = []
    last_ok = 0
    for i, ch in enumerate(text):
        if in_str:
            if escape:
                escape = False
            elif ch == "\\":
                escape = True
            elif ch == '"':
                in_str = False
            continue
        if ch == '"':
            in_str = True
            continue
        if ch in "{[":
            stack.append("}" if ch == "{" else "]")
            last_ok = i
        elif ch in "}]":
            if stack and stack[-1] == ch:
                stack.pop()
                last_ok = i
    if in_str:
        # close string then truncate junk after last complete structure
        text = text + '"'
    if stack:
        # trim to last comma / incomplete key if needed
        text = text.rstrip()
        if text.endswith(","):
            text = text[:-1]
        text = text + "".join(reversed(stack))
    return text


def _extract_json(text: str) -> dict[str, Any]:
    candidates = [_strip_fences(text), _repair_json_text(text)]
    m = re.search(r"\{[\s\S]*\}", text)
    if m:
        candidates.append(m.group(0))
        candidates.append(_repair_json_text(m.group(0)))

    errors: list[str] = []
    for cand in candidates:
        try:
            data = json.loads(cand)
            if isinstance(data, dict):
                return data
        except json.JSONDecodeError as e:
            errors.append(str(e))
    raise json.JSONDecodeError(
        errors[-1] if errors else "Unable to parse LLM JSON",
        text[:200],
        0,
    )


def _compact_solids(solids_geom: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
    """Shrink AABB payload so the model reply is less likely to truncate."""
    out: list[dict[str, Any]] = []
    for s in solids_geom or []:
        center = s.get("center")
        if not center:
            mn, mx = s.get("min"), s.get("max")
            if mn and mx and len(mn) == 3 and len(mx) == 3:
                center = [(mn[i] + mx[i]) / 2 for i in range(3)]
        if not center:
            continue
        out.append(
            {
                "name": s.get("name") or s.get("id"),
                "center": [round(float(x), 1) for x in center[:3]],
            }
        )
    return out


def _draft_from_llm(data: dict[str, Any], fallback: RobotDraft) -> RobotDraft:
    links = [
        LinkDef(name=l["name"], part_names=list(l.get("part_names") or []))
        for l in data.get("links") or []
    ]
    joints = []
    for j in data.get("joints") or []:
        jtype = j.get("joint_type") or "revolute"
        if jtype not in {"revolute", "prismatic", "fixed", "continuous"}:
            jtype = "revolute"
        joints.append(
            JointDef(
                name=j.get("name") or "joint",
                joint_type=jtype,  # type: ignore[arg-type]
                parent=j["parent"],
                child=j["child"],
                origin_xyz=[float(x) for x in (j.get("origin_xyz") or [0, 0, 0])],
                origin_rpy=[float(x) for x in (j.get("origin_rpy") or [0, 0, 0])],
                axis=[float(x) for x in (j.get("axis") or [0, 0, 1])],
                lower=float(j.get("lower", -3.14)),
                upper=float(j.get("upper", 3.14)),
                effort=float(j.get("effort", 100)),
                velocity=float(j.get("velocity", 1)),
                rationale=str(j.get("rationale") or "")[:80],
            )
        )
    if not links:
        return fallback.model_copy(
            update={
                "llm_raw": data,
                "notes": fallback.notes + ["LLM returned no links; kept seed."],
            }
        )
    return RobotDraft(
        name=str(data.get("robot_name") or fallback.name),
        profile=fallback.profile,
        links=links,
        joints=joints,
        notes=list(data.get("notes") or []),
        raw_parts=fallback.raw_parts,
        axis_candidates=fallback.axis_candidates,
        llm_raw=data,
    )


async def _chat(messages: list[dict[str, str]], *, temperature: float = 0.2) -> str:
    headers = {
        "Authorization": f"Bearer {settings.zhipu_api_key}",
        "Content-Type": "application/json",
    }
    body = {
        "model": settings.zhipu_model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": 8192,
    }
    url = f"{settings.zhipu_base_url.rstrip('/')}/chat/completions"
    async with httpx.AsyncClient(timeout=180.0) as client:
        resp = await client.post(url, headers=headers, json=body)
        resp.raise_for_status()
        data = resp.json()
    return data["choices"][0]["message"]["content"]


async def propose_joints_with_zhipu(
    seed: RobotDraft,
    *,
    extra_hint: str = "",
    stats: dict[str, Any] | None = None,
    solids_geom: list[dict[str, Any]] | None = None,
) -> RobotDraft:
    if not settings.zhipu_api_key:
        seed.notes = list(seed.notes) + [
            "ZHIPU_API_KEY 未配置：无法生成关节，请配置后重试。"
        ]
        return seed

    compact = _compact_solids(solids_geom)
    payload_user = {
        "profile": seed.profile,
        "robot_name": seed.name,
        "part_names": seed.raw_parts,
        "solids_center_mm": compact,
        "extra_hint": extra_hint,
        "instruction": (
            "直接生成完整 links+joints JSON。"
            "origin_xyz 用毫米；第一个关节靠近零件 center。"
        ),
    }

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": json.dumps(payload_user, ensure_ascii=False)},
    ]

    content = await _chat(messages, temperature=0.2)
    try:
        parsed = _extract_json(content)
    except json.JSONDecodeError:
        # One repair pass: ask model to fix truncated / invalid JSON
        repair_messages = [
            {"role": "system", "content": "只输出修复后的合法 JSON 对象，不要解释。"},
            {
                "role": "user",
                "content": (
                    "下面内容不是合法 JSON（可能被截断或有尾逗号）。"
                    "请修复并只输出完整 JSON：\n\n" + content[:12000]
                ),
            },
        ]
        repaired = await _chat(repair_messages, temperature=0.0)
        try:
            parsed = _extract_json(repaired)
        except json.JSONDecodeError as e:
            raise ValueError(
                f"智谱返回的 JSON 无法解析（常见于输出被截断）。请重试或减少零件数量。原始错误: {e.msg}"
            ) from e

    return _draft_from_llm(parsed, seed)