# Chat API — loads prompt template, injects dynamic context, calls LLM

import json
import logging
import os
import re
from pathlib import Path

import yaml
from fastapi import APIRouter
from pydantic import BaseModel

from server.config import settings
from server.core.llm_provider.router import router as llm_router

logger = logging.getLogger("testflow")
router = APIRouter()


class ChatRequest(BaseModel):
    message: str
    context: dict = {}
    history: list[dict] = []

# ============================================================================
# Prompt template
# ============================================================================

PROMPT_PATH = Path(__file__).parent.parent.parent / "services" / "prompts" / "chat_system.md"
PROMPT_TEMPLATE = PROMPT_PATH.read_text(encoding="utf-8") if PROMPT_PATH.exists() else "{persona}\n{principles}\n{page_context}\n{data_context}\n{action_format}"


# ============================================================================
# Personality & principles
# ============================================================================

def _load_persona() -> str:
    parts = []
    try:
        p_path = settings.SOUL_DIR / "persona.yaml"
        if p_path.exists():
            data = yaml.safe_load(p_path.open(encoding="utf-8")) or {}
            p = data.get("persona", {})
            parts.append(f"## 你的身份\n你是 {p.get('name','Zeeker')}，一名{p.get('role','资深测试架构师')}。")
            parts.append(f"经验: {p.get('experience','')}")
            parts.append(f"语气: {p.get('tone','专业、严谨、耐心')}")
            traits = p.get("traits", [])
            if traits: parts.append("特点: " + "; ".join(traits))
            limits = p.get("limitations", [])
            if limits: parts.append("限制: " + "; ".join(limits))
    except Exception as e:
        logger.warning("persona load failed: %s", e)
    return "\n".join(parts)


def _load_principles() -> str:
    parts = ["## 不可违背的原则"]
    try:
        p_path = settings.SOUL_DIR / "principles.yaml"
        if p_path.exists():
            data = yaml.safe_load(p_path.open(encoding="utf-8")) or {}
            for pr in data.get("principles", []):
                parts.append(f"- [{pr.get('priority','?')}] {pr.get('rule','')}")
    except Exception as e:
        logger.warning("principles load failed: %s", e)
    return "\n".join(parts)


# ============================================================================
# Dynamic context builders (page-aware)
# ============================================================================

def _build_page_context(ctx: dict) -> str:
    """Describe what the user is looking at on the left panel."""
    page = ctx.get("current_page", "unknown")
    lines = []

    if page == "home":
        lines.append("- 用户正在主页，可以上传 PRD 文件")
        if ctx.get("has_prd") and not ctx.get("has_ir"):
            lines.append("- PRD 已上传，IR 尚未生成。建议用户点击生成 IR。")
        elif ctx.get("has_prd") and ctx.get("has_ir"):
            lines.append("- PRD 和 IR 已生成。用户可以前往 IR 确认或查看测试用例。")

    elif page == "ir-confirm":
        n = len(ctx.get("ir_rules_summary", "").split("\n")) - 2  # skip header rows
        lines.append(f"- 用户正在 **IR 确认页**，当前 IR 含 {max(0,n)} 条规则")
        if ctx.get("ir_id"):
            lines.append(f"- IR ID: {ctx['ir_id']}")
        lines.append("- 用户可以查看思维导图、编辑 YAML、点击节点定位代码")
        lines.append("- 你可以帮用户审查规则完整性、增删改规则、调整优先级")

    elif page == "cases":
        lines.append(f"- 用户正在 **测试用例页**，当前 {ctx.get('testcase_count', 0)} 条用例")
        if ctx.get("tc_set_id"):
            lines.append(f"- 用例集 ID: {ctx['tc_set_id']}")
        if ctx.get("ir_id"):
            lines.append(f"- 这些用例基于 IR: {ctx['ir_id']}")
        lines.append("- 你可以帮用户增加测试用例、检查覆盖率、修改用例内容、导出用例")

        existing = ctx.get("existing_tc_ids", [])
        if existing:
            lines.append(f"\n已有用例 ID: {', '.join(existing[:20])}" + (f" ...共{len(existing)}条" if len(existing) > 20 else ""))

        sel = ctx.get("selected_case")
        if sel and isinstance(sel, dict) and sel.get("id"):
            lines.append(f"\n⚠️ **用户当前选中的测试用例:**")
            lines.append(f"- ID: {sel['id']}  |  标题: {sel.get('title','')}")
            lines.append(f"- 步骤: {sel.get('steps','')[:200]}")
            lines.append(f"- 预期: {sel.get('expected','')}")
            lines.append(f"用户说\"这条\"/\"这个\"时，指的是测试用例 {sel['id']}，不是 IR 规则。在 cases 页面操作的是测试用例！")

    return "\n".join(lines)


def _build_data_context(ctx: dict) -> str:
    """Inject the actual data from the left panel."""
    parts = []

    if ctx.get("prd_text"):
        parts.append(f"### PRD 文本（前 2000 字符）\n```\n{ctx['prd_text'][:2000]}\n```")

    if ctx.get("ir_rules_summary"):
        parts.append(f"### IR 规则列表\n{ctx['ir_rules_summary']}")

    if ctx.get("testcase_count", 0) > 0 and ctx.get("testcase_summary"):
        parts.append(f"### 测试用例概览 ({ctx['testcase_count']} 条)\n```\n{ctx['testcase_summary']}\n```")

    return "\n".join(parts) if parts else "（暂无左侧数据）"


def _build_action_format(ctx: dict) -> str:
    """Page-specific action instructions."""
    page = ctx.get("current_page", "unknown")
    lines = ["修改 IR 规则时使用 rule_id 精准定位 (格式 IR-{FEATURE}-{section}-{NNN})：",
             '  {"action": "add_rule", "rule": {...}}',
             '  {"action": "delete_rule", "rule_id": "IR-ZWM-4.2.1-001"}',
             '  {"action": "modify_rule", "rule_id": "IR-ZWM-4.2.1-001", "changes": {"priority": "P0"}}',
             '  {"action": "generate_cases", "ir_id": "<ir_id>"}']

    if page == "cases":
        lines.append("\n在测试用例页面时，使用 add_case/delete_case/modify_case：")
        lines.append('  {"action": "add_case", "rule": {"case_title":"...","steps":"Given..When..Then..","expected_result":"...","priority":"P0|P1|P2","tags":["正向"],"module":"模块名"}}')
        lines.append('  {"action": "modify_case", "case_id": "TC-ZEEKER-002", "changes": {"case_title": "新标题"}}')
        lines.append('  {"action": "delete_case", "case_id": "TC-ZEEKER-002"}')
        lines.append("注意: add_case 时不要填写 id 字段（系统自动分配）。module 字段尽量填写。delete_case 的 case_id 必须是页面上已有的真实 ID。")
        lines.append("用户要求增改用例时，必须输出 action！只描述不输出 action = 什么都没做。")

    lines.append("\n用户每次要求增删改时，必须在回复末尾输出对应 action 块。不输出 = 没执行。")
    return "\n".join(lines)


# ============================================================================
# Memory
# ============================================================================

def _load_memory(session_id: str) -> list[dict]:
    f = settings.MEMORY_DIR / f"chat_{session_id}.json"
    if f.exists():
        try: return json.loads(f.read_text(encoding="utf-8"))
        except: pass
    return []

def _save_memory(session_id: str, messages: list[dict]):
    os.makedirs(settings.MEMORY_DIR, exist_ok=True)
    f = settings.MEMORY_DIR / f"chat_{session_id}.json"
    try:
        to_save = messages[-30:] if len(messages) > 30 else messages
        f.write_text(json.dumps(to_save, ensure_ascii=False, indent=2), encoding="utf-8")
    except Exception as e:
        logger.warning("memory save failed: %s", e)


# ============================================================================
# Build full system prompt
# ============================================================================

def _build_system_prompt(context: dict) -> str:
    return PROMPT_TEMPLATE.replace("{persona}", _load_persona()) \
                          .replace("{principles}", _load_principles()) \
                          .replace("{page_context}", _build_page_context(context)) \
                          .replace("{data_context}", _build_data_context(context)) \
                          .replace("{action_format}", _build_action_format(context))


# ============================================================================
# Chat endpoint
# ============================================================================

@router.post("")
@router.post("/")
async def chat(request: ChatRequest):
    session_id = request.context.get("session_id", "default")
    memory = _load_memory(session_id)

    system_msg = _build_system_prompt(request.context)
    messages = [{"role": "system", "content": system_msg}]
    messages.extend(memory[-8:])
    for h in (request.history or [])[-6:]:
        messages.append(h)
    messages.append({"role": "user", "content": request.message})

    logger.info("[CHAT] session=%s page=%s memory=%d prompt=%d chars",
                 session_id, request.context.get("current_page","?"),
                 len(memory)//2, len(system_msg))

    client = llm_router.get_chat_client()
    try:
        raw = client.chat(model=llm_router.text_model, messages=messages, temperature=0.3)
    except Exception as e:
        logger.error("[CHAT] LLM failed: %s", e)
        return {"reply": f"LLM 调用失败: {e}", "actions": []}

    reply, actions = _extract_actions(raw)
    memory.append({"role": "user", "content": request.message})
    memory.append({"role": "assistant", "content": reply})
    _save_memory(session_id, memory)

    logger.info("[CHAT] reply=%d chars actions=%d", len(reply), len(actions))
    return {"reply": reply, "actions": actions}


def _extract_actions(text: str) -> tuple[str, list[dict]]:
    actions, clean = [], text
    for m in re.findall(r"```action\s*\n([\s\S]*?)```", text):
        for line in m.strip().split("\n"):
            try:
                a = json.loads(line.strip())
                if isinstance(a, dict) and "action" in a: actions.append(a)
            except: pass
        clean = re.sub(r"```action\s*\n[\s\S]*?```", "", clean)
    for s in re.findall(r'\{"action":\s*"[^"]+"[^}]*\}', text):
        try:
            a = json.loads(s)
            if a not in actions: actions.append(a)
        except: pass
    return clean.strip(), actions