"""ESPD AI chat with tools (doc edit, common fields, RAG, GOST check)."""

from __future__ import annotations

import json
from typing import Any

from app.services.espd.html_util import html_to_plain, plain_to_html
from app.services.espd.models import CONTENT_CONTEXT_KEYS, DOC_TYPE_CONFIG, DocumentType
from app.services.llm import (
    build_context_usage,
    openrouter_chat,
    serialize_tool_calls,
    _extract_json,
)

ESPD_SYSTEM = """Ты помощник по оформлению документов ЕСПД (спецификация, руководства, ПМИ и др.).
Используй tools для предложений правок. Не применяй изменения сам.
Если нужно уточнение — ask_user.
Для проверки по ГОСТам используй gost_check / rag_search.
HTML документа редактируй аккуратно, сохраняя структуру заголовков.
"""

ESPD_TOOLS: list[dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": "edit_doc_html",
            "description": "Заменить HTML-содержимое выбранного типа документа ЕСПД.",
            "parameters": {
                "type": "object",
                "properties": {
                    "doc_type": {
                        "type": "string",
                        "enum": [dt.value for dt in DocumentType if dt != DocumentType.SPEC],
                    },
                    "html": {"type": "string"},
                    "mode": {
                        "type": "string",
                        "enum": ["replace", "append"],
                        "description": "replace — полный текст; append — дописать в конец",
                    },
                },
                "required": ["doc_type", "html"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "set_common_fields",
            "description": "Изменить общие поля основной надписи ЕСПД.",
            "parameters": {
                "type": "object",
                "properties": {
                    "decimal": {"type": "string"},
                    "progname": {"type": "string"},
                    "author": {"type": "string"},
                    "checker": {"type": "string"},
                    "ncontrol": {"type": "string"},
                    "approver": {"type": "string"},
                    "year": {"type": "string"},
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "rag_search",
            "description": "Поиск по базе знаний (глобальные ГОСТы + документы проекта).",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string"},
                    "top_k": {"type": "integer"},
                    "gost_only": {"type": "boolean"},
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "gost_check",
            "description": "Проверить текст/документ на соответствие нормам из RAG (с цитатами).",
            "parameters": {
                "type": "object",
                "properties": {
                    "doc_type": {"type": "string"},
                    "focus": {"type": "string", "description": "На что обратить внимание"},
                    "use_current_html": {"type": "boolean"},
                    "text": {"type": "string"},
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "ask_user",
            "description": "Уточняющий вопрос без правок.",
            "parameters": {
                "type": "object",
                "properties": {"question": {"type": "string"}},
                "required": ["question"],
            },
        },
    },
]


def _parse_args(raw: Any) -> dict[str, Any]:
    if isinstance(raw, dict):
        return raw
    try:
        data = json.loads(raw or "{}")
        return data if isinstance(data, dict) else {}
    except json.JSONDecodeError:
        return {}


def espd_tool_calls_to_edits(
    tool_calls: list[dict[str, Any]],
) -> tuple[list[dict[str, Any]], str | None, list[dict[str, Any]]]:
    """Returns (proposal_edits, ask_question, executable_tool_requests)."""
    edits: list[dict[str, Any]] = []
    exec_reqs: list[dict[str, Any]] = []
    question: str | None = None
    for tc in tool_calls:
        fn = tc.get("function") or {}
        name = fn.get("name") or ""
        args = _parse_args(fn.get("arguments"))
        if name == "ask_user":
            question = str(args.get("question") or "").strip() or question
        elif name == "edit_doc_html":
            edits.append(
                {
                    "op": "edit_doc_html",
                    "doc_type": args.get("doc_type"),
                    "html": args.get("html") or "",
                    "mode": args.get("mode") or "replace",
                }
            )
        elif name == "set_common_fields":
            fields = {k: v for k, v in args.items() if k in {
                "decimal", "progname", "author", "checker", "ncontrol", "approver", "year"
            } and v is not None}
            if fields:
                edits.append({"op": "set_common_fields", "fields": fields})
        elif name in ("rag_search", "gost_check"):
            exec_reqs.append({"name": name, "arguments": args, "tool_call_id": tc.get("id")})
    return edits, question, exec_reqs


async def chat_espd(
    message: str,
    *,
    history: list[dict[str, str]] | None = None,
    common: dict[str, Any] | None = None,
    doc_contents: dict[str, str] | None = None,
    active_doc_type: str | None = None,
    project_id: int | None = None,
    rag_search_fn=None,
    gost_check_fn=None,
) -> dict[str, Any]:
    """ESPD assistant with multi-round tools. rag_search_fn/gost_check_fn optional callables."""
    hist = list(history or [])
    if not hist or hist[-1].get("content") != message:
        hist.append({"role": "user", "content": message})

    docs_summary = {
        k: (html_to_plain(v)[:2000] if v else "") for k, v in (doc_contents or {}).items()
    }
    catalog = [
        {"id": dt.value, "name": cfg["name"], "content_key": CONTENT_CONTEXT_KEYS.get(dt, "")}
        for dt, cfg in DOC_TYPE_CONFIG.items()
        if dt != DocumentType.SPEC
    ]
    snapshot = (
        f"Активный документ: {active_doc_type or '—'}\n"
        f"Общие поля: {json.dumps(common or {}, ensure_ascii=False)}\n"
        f"Типы документов: {json.dumps(catalog, ensure_ascii=False)}\n"
        f"Тексты (plain, усечённо): {json.dumps(docs_summary, ensure_ascii=False)}"
    )
    messages: list[dict[str, Any]] = [
        {"role": "system", "content": ESPD_SYSTEM},
        {"role": "user", "content": snapshot},
        {"role": "assistant", "content": "Контекст ЕСПД принят. Жду команды."},
    ]
    for m in hist[-20:]:
        if m.get("content"):
            messages.append({"role": m["role"], "content": m["content"]})

    all_edits: list[dict[str, Any]] = []
    all_tools: list[dict[str, Any]] = []
    question: str | None = None
    final_content = ""
    rag_notes: list[str] = []

    for round_no in range(1, 5):
        msg = await openrouter_chat(
            messages,
            temperature=0.2,
            max_tokens=3500,
            tools=ESPD_TOOLS,
            tool_choice="auto",
        )
        content = (msg.get("content") or "").strip()
        tool_calls = msg.get("tool_calls") or []
        if tool_calls:
            all_tools.extend(serialize_tool_calls(tool_calls, round_no=round_no))
            edits, q, exec_reqs = espd_tool_calls_to_edits(tool_calls)
            all_edits.extend(edits)
            if q:
                question = q
                final_content = content
                break

            messages.append({"role": "assistant", "content": content or "", "tool_calls": tool_calls})
            for req in exec_reqs:
                result_payload: dict[str, Any]
                if req["name"] == "rag_search" and rag_search_fn:
                    result_payload = await _maybe_await(
                        rag_search_fn,
                        req["arguments"].get("query") or "",
                        int(req["arguments"].get("top_k") or 5),
                        bool(req["arguments"].get("gost_only")),
                        project_id,
                    )
                    rag_notes.append(json.dumps(result_payload, ensure_ascii=False)[:4000])
                elif req["name"] == "gost_check" and gost_check_fn:
                    html = ""
                    dt = req["arguments"].get("doc_type") or active_doc_type
                    if req["arguments"].get("use_current_html", True) and dt:
                        html = (doc_contents or {}).get(dt, "")
                    text = req["arguments"].get("text") or html_to_plain(html)
                    result_payload = await _maybe_await(
                        gost_check_fn,
                        text,
                        req["arguments"].get("focus") or "",
                        project_id,
                    )
                    rag_notes.append(json.dumps(result_payload, ensure_ascii=False)[:6000])
                    if result_payload.get("reply"):
                        final_content = str(result_payload["reply"])
                else:
                    result_payload = {
                        "status": "accepted_as_proposal" if req["name"] not in ("rag_search", "gost_check") else "unavailable",
                        "tool": req["name"],
                    }
                # also acknowledge propose tools
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": req.get("tool_call_id") or f"call_{round_no}_{req['name']}",
                        "content": json.dumps(result_payload, ensure_ascii=False),
                    }
                )
            # acknowledge edit proposals without exec
            for tc in tool_calls:
                fn = (tc.get("function") or {}).get("name")
                if fn in ("edit_doc_html", "set_common_fields"):
                    messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": tc.get("id") or f"call_{round_no}_{fn}",
                            "content": json.dumps(
                                {"status": "accepted_as_proposal", "tool": fn},
                                ensure_ascii=False,
                            ),
                        }
                    )
            if exec_reqs and not edits:
                continue
            if edits and not exec_reqs:
                # one more round for natural reply optional
                continue
            continue

        if content:
            parsed = _extract_json(content)
            if isinstance(parsed, dict) and parsed.get("reply"):
                final_content = str(parsed["reply"])
            else:
                final_content = content
        break

    reply = question or final_content
    if not reply:
        if all_edits:
            reply = f"Предложено правок: {len(all_edits)}. Подтвердите применение."
        elif rag_notes:
            reply = "Результаты поиска по базе знаний получены (см. инструменты)."
        else:
            reply = "Нет предложенных изменений."

    return {
        "reply": reply,
        "edits": all_edits,
        "tools": all_tools,
        "context": build_context_usage(messages),
    }


async def _maybe_await(fn, *args):
    result = fn(*args)
    if hasattr(result, "__await__"):
        return await result
    return result


def apply_espd_edits(state, edits: list[dict[str, Any]]):
    """Apply proposal edits onto EspdState (mutates)."""
    for edit in edits:
        op = edit.get("op")
        if op == "edit_doc_html":
            dt = edit.get("doc_type")
            if not dt:
                continue
            html = edit.get("html") or ""
            mode = edit.get("mode") or "replace"
            if mode == "append":
                prev = state.get_doc_html(dt)
                html = (prev or "") + html
            state.set_doc_html(dt, html)
        elif op == "set_common_fields":
            fields = edit.get("fields") or {}
            for k, v in fields.items():
                if hasattr(state.common, k) and v is not None:
                    setattr(state.common, k, str(v))
    # sync legacy
    if state.get_doc_html("ro"):
        state.ro_ai_content = html_to_plain(state.get_doc_html("ro"))
    return state