"""
Studio worldbook step LLM responder (R3/R4).

Assembles R2 context blocks + short step dialogue, calls LLM for structured JSON,
returns thinking, draft, questions, and evaluation.
"""
from __future__ import annotations

import json
import re
import uuid
from datetime import datetime
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage

from models.studio_models import (
    LastToolResponse,
    PromptBlock,
    StudioNode,
    StepMessage,
    ToolQuestionOption,
)
from services.studio_context_service import assemble_prompt_blocks
from utils.llm_client import LLMClient

_llm_client = LLMClient()

_JSON_FENCE = re.compile(r"```(?:json)?\s*([\s\S]*?)```", re.IGNORECASE)


def resolve_api_config(
    profile_id: Optional[str],
    api_config: Optional[Dict[str, str]],
) -> Dict[str, str]:
    """Merge frontend apiConfig with stored profile mainLLM key (same as chat WS)."""
    resolved = dict(api_config or {})
    if profile_id:
        try:
            try:
                from api.routes.apiConfigRoute import load_profile
            except ImportError:
                from backend.api.routes.apiConfigRoute import load_profile

            profile = load_profile(profile_id)
            if profile:
                main_llm = profile.get("apis", {}).get("mainLLM", {})
                if main_llm.get("apiUrl") and not resolved.get("api_url"):
                    resolved["api_url"] = main_llm.get("apiUrl", "")
                if main_llm.get("model") and not resolved.get("model"):
                    resolved["model"] = main_llm.get("model", "")
                api_key = main_llm.get("apiKey", "")
                if api_key:
                    resolved["api_key"] = api_key
        except Exception as exc:
            print(f"[StudioStepRespond] 加载 API 配置失败: {exc}")
    return resolved


def _blocks_to_context_text(blocks: List[PromptBlock]) -> str:
    sections: List[str] = []
    for block in blocks:
        sections.append(f"## {block.label}\n{block.content}")
    return "\n\n".join(sections)


def _build_system_prompt(node: StudioNode) -> str:
    insertion = (node.config or {}).get("insertion") or {}
    key = insertion.get("key") or "（未配置关键词）"
    comment = insertion.get("comment") or ""

    return f"""你是 Studio 创作助手，负责为当前流水线步骤生成或修订世界书条目草稿。

当前步骤：{node.displayName}
目标关键词：{key}
备注：{comment or "（无）"}

你必须只输出一个 JSON 对象（不要 markdown 代码块外的其他文字），字段如下：
{{
  "thinking": "你的内部思考过程（逐步推理，中文）",
  "currentProduct": "世界书条目正文（纯文本或 Markdown，可直接写入条目 content）",
  "questions": [
    {{
      "question": "需要用户澄清的问题",
      "options": ["选项A", "选项B", "选项C"]
    }}
  ],
  "evaluation": "对照评价维度的自检与优化建议（中文，面向用户）"
}}

规则：
1. currentProduct 必须是完整、可注入世界书的条目正文。
2. questions 为 0–3 条；每条至少 2 个 options；若无需澄清则 questions 为空数组。
3. evaluation 需引用上下文中的评价标准，给出具体、可操作的反馈。
4. 若用户要求修改，在 currentProduct 中输出修订后的完整条目，而非仅说明改了什么。
5. 全部字段使用中文（专有名词除外）。"""


def _dialogue_to_langchain(
    step_messages: List[StepMessage],
) -> List[Any]:
    messages: List[Any] = []
    for msg in step_messages:
        if msg.role == "user":
            messages.append(HumanMessage(content=msg.content))
        elif msg.role == "assistant":
            messages.append(AIMessage(content=msg.content))
    return messages


def _build_llm_messages(
    node: StudioNode,
    prompt_blocks: List[PromptBlock],
    step_messages: List[StepMessage],
    user_message: str,
) -> List[Any]:
    context_text = _blocks_to_context_text(prompt_blocks)
    system_prompt = _build_system_prompt(node)

    messages: List[Any] = [SystemMessage(content=system_prompt)]
    messages.append(
        HumanMessage(
            content=f"以下为当前步骤上下文（不含完整聊天历史）：\n\n{context_text}"
        )
    )
    messages.extend(_dialogue_to_langchain(step_messages))
    messages.append(HumanMessage(content=user_message))
    return messages


def _validate_api_config(api_config: Dict[str, str]) -> None:
    if not api_config.get("api_key"):
        raise ValueError("API Key 未配置，请先在 API 配置页面保存密钥")
    if not api_config.get("api_url"):
        raise ValueError("API 地址未配置，请先在 API 配置页面保存 mainLLM")


def _extract_json(raw: str) -> Dict[str, Any]:
    text = (raw or "").strip()
    if not text:
        raise ValueError("模型返回为空")

    fence = _JSON_FENCE.search(text)
    if fence:
        text = fence.group(1).strip()

    try:
        return json.loads(text)
    except json.JSONDecodeError:
        start = text.find("{")
        end = text.rfind("}")
        if start >= 0 and end > start:
            return json.loads(text[start : end + 1])
        raise ValueError("无法解析模型返回的 JSON")


def _decode_json_string_partial(raw: str) -> str:
    """Decode a possibly incomplete JSON string body (no surrounding quotes)."""
    out: List[str] = []
    i = 0
    while i < len(raw):
        if raw[i] == "\\" and i + 1 < len(raw):
            nxt = raw[i + 1]
            if nxt == "n":
                out.append("\n")
            elif nxt == "t":
                out.append("\t")
            elif nxt == "r":
                out.append("\r")
            elif nxt == '"':
                out.append('"')
            elif nxt == "\\":
                out.append("\\")
            elif nxt == "/":
                out.append("/")
            elif nxt == "u" and i + 5 < len(raw):
                try:
                    out.append(chr(int(raw[i + 2 : i + 6], 16)))
                    i += 6
                    continue
                except ValueError:
                    out.append(nxt)
            else:
                out.append(nxt)
            i += 2
        else:
            out.append(raw[i])
            i += 1
    return "".join(out)


def _extract_partial_thinking(raw: str) -> Optional[str]:
    """Best-effort extraction of thinking field from incomplete JSON stream."""
    marker = '"thinking"'
    idx = raw.find(marker)
    if idx < 0:
        return None

    colon = raw.find(":", idx + len(marker))
    if colon < 0:
        return None

    rest = raw[colon + 1 :].lstrip()
    if not rest.startswith('"'):
        return None

    body_start = 1
    i = body_start
    while i < len(rest):
        ch = rest[i]
        if ch == '"':
            break
        if ch == "\\":
            i += 2
            continue
        i += 1

    partial = rest[body_start:i]
    if not partial:
        return None
    return _decode_json_string_partial(partial)


def _normalize_draft(
    current_product: Any,
    node: StudioNode,
    existing_draft: Optional[Dict[str, Any]],
) -> Dict[str, Any]:
    draft: Dict[str, Any] = dict(existing_draft or {})
    insertion = (node.config or {}).get("insertion") or {}

    if isinstance(current_product, str):
        draft["entryContent"] = current_product.strip()
    elif isinstance(current_product, dict):
        draft.update(current_product)
        if "entryContent" not in draft and "content" in draft:
            draft["entryContent"] = draft["content"]
    else:
        draft["entryContent"] = str(current_product)

    if insertion.get("key"):
        draft["insertionKey"] = insertion["key"]
    if insertion.get("comment"):
        draft["insertionComment"] = insertion["comment"]
    draft["nodeId"] = node.id
    draft["displayName"] = node.displayName
    return draft


def _normalize_questions(raw: Any) -> List[ToolQuestionOption]:
    if not isinstance(raw, list):
        return []
    result: List[ToolQuestionOption] = []
    for item in raw:
        if not isinstance(item, dict):
            continue
        question = (item.get("question") or "").strip()
        if not question:
            continue
        options = [
            str(o).strip()
            for o in (item.get("options") or [])
            if str(o).strip()
        ]
        if len(options) < 2:
            continue
        result.append(ToolQuestionOption(question=question, options=options))
    return result[:3]


def _assistant_message_text(parsed: Dict[str, Any]) -> str:
    evaluation = (parsed.get("evaluation") or "").strip()
    if evaluation:
        return evaluation
    product = parsed.get("currentProduct")
    if isinstance(product, str) and product.strip():
        preview = product.strip()
        if len(preview) > 400:
            preview = preview[:400] + "…"
        return f"已更新条目草稿：\n\n{preview}"
    return "已处理您的消息，请查看左侧目前产物。"


def _build_turn_result(
    parsed: Dict[str, Any],
    *,
    node: StudioNode,
    user_message: str,
    existing_draft: Optional[Dict[str, Any]],
) -> Tuple[Dict[str, Any], LastToolResponse, StepMessage, StepMessage]:
    now = datetime.now().isoformat()

    last_draft = _normalize_draft(
        parsed.get("currentProduct"),
        node,
        existing_draft,
    )
    last_tool_response = LastToolResponse(
        thinking=(parsed.get("thinking") or "").strip() or None,
        evaluation=(parsed.get("evaluation") or "").strip() or None,
        questions=_normalize_questions(parsed.get("questions")),
        generatedAt=now,
    )

    user_step_msg = StepMessage(
        id=f"msg-{uuid.uuid4().hex[:12]}",
        role="user",
        content=user_message,
        createdAt=now,
    )
    assistant_step_msg = StepMessage(
        id=f"msg-{uuid.uuid4().hex[:12]}",
        role="assistant",
        content=_assistant_message_text(parsed),
        createdAt=now,
    )

    return last_draft, last_tool_response, user_step_msg, assistant_step_msg


async def studio_step_respond(
    *,
    node: StudioNode,
    prompt_blocks: List[PromptBlock],
    step_messages: List[StepMessage],
    user_message: str,
    existing_draft: Optional[Dict[str, Any]],
    api_config: Dict[str, str],
    stream: bool = False,
) -> Tuple[Dict[str, Any], LastToolResponse, StepMessage, StepMessage]:
    """
    Execute one worldbook step turn (non-streaming).

    Returns (last_draft, last_tool_response, user_step_msg, assistant_step_msg).
    """
    _validate_api_config(api_config)
    messages = _build_llm_messages(node, prompt_blocks, step_messages, user_message)
    model = api_config.get("model") or "gpt-4o-mini"

    if stream:
        print("[StudioStepRespond] stream=True 应使用 studio_step_respond_stream")

    response = await _llm_client.chat_completion(
        messages=messages,
        api_url=api_config.get("api_url", ""),
        api_key=api_config.get("api_key", ""),
        model=model,
        temperature=0.7,
        max_tokens=8000,
        request_timeout=120,
        stream=False,
    )

    raw_content = ""
    if isinstance(response, dict):
        raw_content = (
            response.get("choices", [{}])[0]
            .get("message", {})
            .get("content", "")
        )
    else:
        raw_content = str(response)

    parsed = _extract_json(raw_content)
    return _build_turn_result(
        parsed,
        node=node,
        user_message=user_message,
        existing_draft=existing_draft,
    )


async def studio_step_respond_stream(
    *,
    node: StudioNode,
    prompt_blocks: List[PromptBlock],
    step_messages: List[StepMessage],
    user_message: str,
    existing_draft: Optional[Dict[str, Any]],
    api_config: Dict[str, str],
) -> AsyncGenerator[Dict[str, Any], None]:
    """
    Stream thinking field while LLM generates structured JSON (R4).

    Yields:
      - {"type": "thinking_delta", "content": "..."}
      - {"type": "complete", "last_draft", "last_tool_response", "user_msg", "assistant_msg"}
    """
    _validate_api_config(api_config)
    messages = _build_llm_messages(node, prompt_blocks, step_messages, user_message)
    model = api_config.get("model") or "gpt-4o-mini"

    accumulated = ""
    last_thinking = ""

    async for chunk in _llm_client.stream_chat(
        messages=messages,
        api_url=api_config.get("api_url", ""),
        api_key=api_config.get("api_key", ""),
        model=model,
        temperature=0.7,
        max_tokens=8000,
        request_timeout=120,
    ):
        if chunk.get("type") != "chunk":
            continue
        accumulated += chunk.get("content") or ""
        partial = _extract_partial_thinking(accumulated)
        if partial and partial != last_thinking:
            last_thinking = partial
            yield {"type": "thinking_delta", "content": partial}

    parsed = _extract_json(accumulated)
    last_draft, last_tool_response, user_msg, assistant_msg = _build_turn_result(
        parsed,
        node=node,
        user_message=user_message,
        existing_draft=existing_draft,
    )

    yield {
        "type": "complete",
        "last_draft": last_draft,
        "last_tool_response": last_tool_response.model_dump(mode="json"),
        "user_msg": user_msg.model_dump(mode="json"),
        "assistant_msg": assistant_msg.model_dump(mode="json"),
    }