"""Article Annotator - 在原文中标注关键信息提取区域.

将提取结果映射回原文，用 {{path||value||raw_text||constraint||transform}} 格式
标注每个关键信息的提取位置。

核心流程:
    Instance JSON + Definition + Article
        → 扁平化提取值列表 (路径 + 值 + 元数据)
        → [LLM Step 2a] 为每个提取值定位原文中的精确文本片段
        → [LLM Step 2b] 基于精确原文片段生成值约束和转换规则
        → 组装标注 Prompt → [LLM Step 3] 在原文中标注
        → 标注后的文章文本

Step 2a 先定位精确原文片段，Step 2b 基于片段生成规则，
避免 LLM 在整篇文章中识别提取值对应原文时产生边界模糊。
"""

from __future__ import annotations

import json
import logging
from typing import Any

from src.schema_flattener import DefinitionNode

logger = logging.getLogger(__name__)


# ═══════════════════════════════════════════════════════════════════════════════
# 转义工具
# ═══════════════════════════════════════════════════════════════════════════════

_ESCAPE_MAP = {
    "\\": "\\\\",
    "|": "\\|",
    "}": "\\}",
}


def escape_field(text: str) -> str:
    """转义标注字段中的特殊字符.

    转义规则: ``\\`` → ``\\\\``, ``|`` → ``\\|``, ``}`` → ``\\}``
    """
    for char, escaped in _ESCAPE_MAP.items():
        text = text.replace(char, escaped)
    return text


def unescape_field(text: str) -> str:
    """反转义标注字段."""
    result: list[str] = []
    i = 0
    while i < len(text):
        if text[i] == "\\" and i + 1 < len(text):
            result.append(text[i + 1])
            i += 2
        else:
            result.append(text[i])
            i += 1
    return "".join(result)


# ═══════════════════════════════════════════════════════════════════════════════
# 显式约束解析 (来自 definition JSON 中的 constraints 字段)
# ═══════════════════════════════════════════════════════════════════════════════

def parse_raw_constraints(metadata: dict[str, Any]) -> str:
    """解析 definition 中显式定义的约束为自然语言.

    仅处理 metadata.constraints.raw_constraints，
    不做任何推断，推断交给 LLM。
    """
    constraints = metadata.get("constraints", {})
    raw = constraints.get("raw_constraints", {})
    if not raw:
        return ""

    parts: list[str] = []

    # 数值范围
    if "min_value" in raw and "max_value" in raw:
        min_v = _first_element(raw["min_value"])
        max_v = _first_element(raw["max_value"])
        parts.append(f"取值范围: {min_v}~{max_v}")
    elif "min_value" in raw:
        parts.append(f"最小值: {_first_element(raw['min_value'])}")
    elif "max_value" in raw:
        parts.append(f"最大值: {_first_element(raw['max_value'])}")

    # 枚举字符串
    if "enum_string" in raw:
        for enum_list in raw["enum_string"]:
            if isinstance(enum_list, list):
                parts.append(f"枚举值: {'/'.join(str(v) for v in enum_list)}")
            else:
                parts.append(f"枚举值: {enum_list}")

    # 正则模式
    if "pattern" in raw:
        parts.append(f"格式: {raw['pattern']}")

    # 字符串长度
    if "min_length" in raw or "max_length" in raw:
        min_l = raw.get("min_length", 0)
        max_l = raw.get("max_length", "∞")
        parts.append(f"长度范围: {min_l}~{max_l}")

    # 附加原始 JSON 以便 LLM 理解
    parts.append(f"[原始约束: {json.dumps(raw, ensure_ascii=False)}]")

    return "; ".join(parts)


def _first_element(value: Any) -> Any:
    """如果值是列表则取第一个元素，否则原样返回."""
    if isinstance(value, list) and value:
        return value[0]
    return value


# ═══════════════════════════════════════════════════════════════════════════════
# 实例扁平化
# ═══════════════════════════════════════════════════════════════════════════════

def flatten_instance_for_annotation(
    definition: DefinitionNode,
    instance: Any,
    parent_path: str = "",
) -> list[dict[str, Any]]:
    """将实例 JSON 扁平化为标注项列表.

    并行遍历 definition 和 instance，为每个叶子属性生成:
    - path: 完整节点路径 (含列表索引)
    - value: 提取值
    - metadata: 节点元数据 (含 node_name, value_type 等)
    - explicit_constraint: 显式约束文本 (来自 definition)
    - constraint: 值约束 (初始为空，由 LLM 填充)
    - transform: 转换规则 (初始为空，由 LLM 填充)

    Returns:
        标注项列表，每项对应一个叶子属性
    """
    metadata = definition.get("metadata", {})
    label = metadata.get("node_label", "")
    node_type = metadata.get("node_type", "")
    value_type = metadata.get("value_type", "")
    is_list = metadata.get("is_list", False)

    current_path = f"{parent_path}.{label}" if parent_path else label
    items: list[dict[str, Any]] = []

    if node_type == "concept":
        _flatten_concept(definition, instance, current_path, items)
    elif value_type == "class":
        _flatten_class(definition, instance, current_path, is_list, metadata, items)
    else:
        # 叶子属性 (基础类型)
        if is_list and isinstance(instance, list):
            for i, element in enumerate(instance):
                items.append(_make_annotation_item(
                    f"{current_path}[{i}]", element, metadata,
                ))
        else:
            items.append(_make_annotation_item(current_path, instance, metadata))

    return items


def _flatten_concept(
    definition: DefinitionNode,
    instance: Any,
    current_path: str,
    items: list[dict[str, Any]],
) -> None:
    """扁平化 concept 节点."""
    if not isinstance(instance, dict):
        return
    for key, subdef in definition.get("subnodes", {}).items():
        if key in instance:
            items.extend(flatten_instance_for_annotation(
                subdef, instance[key], current_path,
            ))


def _flatten_class(
    definition: DefinitionNode,
    instance: Any,
    current_path: str,
    is_list: bool,
    metadata: dict[str, Any],
    items: list[dict[str, Any]],
) -> None:
    """扁平化 class 类型节点."""
    subnodes = definition.get("subnodes", {})

    if is_list:
        if not isinstance(instance, list):
            return
        for i, element in enumerate(instance):
            indexed_path = f"{current_path}[{i}]"
            if isinstance(element, dict):
                for subkey, subdef in subnodes.items():
                    if subkey in element:
                        items.extend(flatten_instance_for_annotation(
                            subdef, element[subkey], indexed_path,
                        ))
    else:
        if not isinstance(instance, dict):
            return
        for subkey, subdef in subnodes.items():
            if subkey in instance:
                items.extend(flatten_instance_for_annotation(
                    subdef, instance[subkey], current_path,
                ))


def _make_annotation_item(
    path: str,
    value: Any,
    metadata: dict[str, Any],
) -> dict[str, Any]:
    """构造单个标注项.

    constraint 和 transform 初始为空，后续由 LLM 辅助填充。
    """
    return {
        "path": path,
        "value": value,
        "metadata": metadata,
        "explicit_constraint": parse_raw_constraints(metadata),
        "constraint": "",
        "transform": "",
    }


# ═══════════════════════════════════════════════════════════════════════════════
# LLM 辅助 Step 2a：定位提取值对应的原文片段
# ═══════════════════════════════════════════════════════════════════════════════

def build_raw_text_prompt(
    article: str,
    items: list[dict[str, Any]],
) -> list[dict[str, str]]:
    """构建让 LLM 为每个提取值定位原文片段的 Prompt.

    轻量级调用，不需要推理过程，快速返回每个提取值在原文中对应的
    精确文本片段。

    Returns:
        消息列表 [{"role": ..., "content": ...}, ...]
    """
    items_desc: list[str] = []
    for i, item in enumerate(items, 1):
        value = item["value"]
        value_str = json.dumps(value, ensure_ascii=False) if not isinstance(value, str) else value
        meta = item["metadata"]
        node_name = meta.get("node_name", "")

        line = f"{i}. 路径: {item['path']}"
        line += f"\n   字段名: {node_name}"
        line += f"\n   提取值: {value_str}"
        items_desc.append(line)

    items_text = "\n\n".join(items_desc)

    system = (
        "你是一个文本定位助手。你的任务是在文章中找到每个提取值对应的最精确的原文片段。\n\n"
        "## 规则\n\n"
        "1. **精确边界**：只包含提取值本身对应的文字，不要包含无关的上下文。"
        "例如提取值是 3000.0，原文是\"3000元\"，应该返回\"3000\"而不是\"3000元\"。\n"
        "2. **完整映射**：如果提取值是原文多个片段的合成，返回所有相关片段，用 | 分隔。\n"
        "3. **格式转换**：提取值可能是格式转换后的结果（如日期转 ISO、布尔转 true/false），"
        "请找到转换前的原文表述。\n"
        "4. **完整提取**：如果某个提取值在原文中拥有多个对应文本，如果对应文本表述存在差别，"
        "对这些文本分别进行提取，并处理成多个对象。\n"
        "5. **找不到时**：如果在原文中确实找不到对应文本，返回空字符串。\n\n"
        "严格按以下 JSON 格式输出，不要输出其他内容：\n"
        "```json\n"
        "[\n"
        '  {"path": "字段路径", "raw_text": "原文中对应的精确文本片段"},\n'
        "  ...\n"
        "]\n"
        "```"
    )

    user = (
        f"## 文章内容\n\n<article>\n{article}\n</article>\n\n"
        f"## 需要定位原文的提取值\n\n{items_text}\n\n"
        "请在文章中找到每个提取值对应的精确原文片段。只输出 JSON 数组。"
    )

    return [
        {"role": "system", "content": system},
        {"role": "user", "content": user},
    ]


def parse_raw_text_response(
    response: str,
    items: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """解析 LLM 返回的原文定位 JSON，填充到 items 的 raw_text 字段."""
    import re

    # 提取 JSON
    json_match = re.search(r"```json\s*\n?(.*?)\n?\s*```", response, re.DOTALL)
    if json_match:
        json_str = json_match.group(1)
    else:
        json_match = re.search(r"\[[\s\S]*\]", response, re.DOTALL)
        json_str = json_match.group(0) if json_match else None

    if not json_str:
        logger.warning("原文定位解析失败，使用空值")
        return items

    try:
        results = json.loads(json_str.strip())
    except json.JSONDecodeError:
        logger.warning("原文定位 JSON 解析失败，使用空值")
        return items

    if not isinstance(results, list):
        return items

    # 建立 path → raw_text 的映射
    raw_map: dict[str, str] = {}
    for entry in results:
        if isinstance(entry, dict) and "path" in entry:
            raw_map[entry["path"]] = entry.get("raw_text", "")

    # 填充到 items
    for item in items:
        raw_text = raw_map.get(item["path"], "")
        item["raw_text"] = raw_text

    return items


# ═══════════════════════════════════════════════════════════════════════════════
# LLM 辅助 Step 2b：约束和转换规则生成（基于精确原文片段）
# ═══════════════════════════════════════════════════════════════════════════════

def build_constraint_transform_prompt(
    items: list[dict[str, Any]],
) -> list[dict[str, str]]:
    """构建让 LLM 为每个字段生成约束和转换规则的 Prompt.

    基于精确的原文片段（由 Step 2a 定位），不再传入整篇文章。
    轻量级调用，不需要推理过程（temperature=0），快速返回。

    Args:
        items: 标注项列表，每项需包含 path, value, metadata, raw_text 字段

    Returns:
        消息列表 [{"role": ..., "content": ...}, ...]
    """
    items_desc: list[str] = []
    for i, item in enumerate(items, 1):
        value = item["value"]
        value_str = json.dumps(value, ensure_ascii=False) if not isinstance(value, str) else value
        meta = item["metadata"]
        node_name = meta.get("node_name", "")
        value_type = meta.get("value_type", "")
        explicit = item.get("explicit_constraint", "")
        raw_text = item.get("raw_text", "")

        line = f"{i}. 路径: {item['path']}"
        line += f"\n   字段名: {node_name}"
        line += f"\n   值类型: {value_type}"
        line += f"\n   提取值: {value_str}"
        line += f"\n   原文片段: {raw_text}"
        if explicit:
            line += f"\n   显式约束: {explicit}"
        items_desc.append(line)

    items_text = "\n\n".join(items_desc)

    system = (
        "你是一个约束分析和转换规则生成助手。"
        "根据字段的提取值和对应的原文片段，为每个字段生成两条信息：\n"
        "1. **值约束**：描述该字段取值的限制条件。如果定义中有显式约束，"
        "直接使用显式约束；否则根据字段名、值类型推断合理的约束；"
        "无法确定的不给出。\n"
        "2. **转换规则**：描述提取值如何映射回原文片段。"
        "注意：由于提取值可能发生变化，转换规则应全面，例如时间的转换规则，若原文中含有\"凌晨/早上/早/上午/中午/下午/晚上/晚\"等其中之一，你就应该考虑全面，将其他情况的规则也梳理出来。"
        "例如日期字段的格式映射(ISO 8601 → YYYY年M月D日(凌晨/早上/早/上午/中午/下午/晚上/晚)H点[m分[s秒]]（1.根据是否为整点，可视情省略[]中的内容；2.0点-6点为凌晨，6点-9点为早/早上，9点到11点为上午，11点到13点为中午，13点到18点为下午，18点到0点为晚/晚上）)、"
        "布尔字段的语义映射(true->已完成备案/false->未完成备案)、"
        "数值字段描述提取值与原文数值的对应关系（如 3000.0 → 3000）、"
        "字符串的调整，如果提取值与原文片段完全一致，填'不作转换'；如果从提取值可以推理出原文片段，填写对应的推理规则；否则，填'根据原文进行微调'。\n\n"
        "严格按以下 JSON 格式输出，不要输出其他内容：\n"
        "```json\n"
        "[\n"
        '  {"path": "字段路径", "constraint": "值约束", "transform": "转换规则"},\n'
        "  ...\n"
        "]\n"
        "```"
    )

    user = (
        f"## 需要生成约束和转换规则的字段\n\n{items_text}\n\n"
        "请根据每个字段的提取值和原文片段，生成 constraint 和 transform。只输出 JSON 数组。"
    )

    return [
        {"role": "system", "content": system},
        {"role": "user", "content": user},
    ]


def parse_constraint_transform_response(
    response: str,
    items: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """解析 LLM 返回的约束和转换规则 JSON，填充到 items 中."""
    import re

    # 提取 JSON
    json_match = re.search(r"```json\s*\n?(.*?)\n?\s*```", response, re.DOTALL)
    if json_match:
        json_str = json_match.group(1)
    else:
        json_match = re.search(r"\[[\s\S]*\]", response, re.DOTALL)
        json_str = json_match.group(0) if json_match else None

    if not json_str:
        logger.warning("约束/转换规则解析失败，使用空值")
        return items

    try:
        rules = json.loads(json_str.strip())
    except json.JSONDecodeError:
        logger.warning("约束/转换规则 JSON 解析失败，使用空值")
        return items

    if not isinstance(rules, list):
        return items

    # 建立 path → rule 的映射
    rule_map: dict[str, dict[str, str]] = {}
    for rule in rules:
        if isinstance(rule, dict) and "path" in rule:
            rule_map[rule["path"]] = rule

    # 填充到 items
    for item in items:
        rule = rule_map.get(item["path"])
        if rule:
            item["constraint"] = rule.get("constraint", "")
            item["transform"] = rule.get("transform", "")

    return items


# ═══════════════════════════════════════════════════════════════════════════════
# 标注格式化
# ═══════════════════════════════════════════════════════════════════════════════

def format_annotation_block(
    path: str,
    value: Any,
    raw_text: str,
    constraint: str,
    transform: str,
) -> str:
    """将标注数据格式化为 {{path||value||raw||constraint||transform}} 块."""
    value_str = json.dumps(value, ensure_ascii=False) if not isinstance(value, str) else value
    parts = [
        escape_field(path),
        escape_field(value_str),
        escape_field(raw_text),
        escape_field(constraint),
        escape_field(transform),
    ]
    return "{{" + "||".join(parts) + "}}"


def format_annotation_items_text(items: list[dict[str, Any]]) -> str:
    """将标注项列表格式化为 LLM 可读的文本.

    每个标注项末尾附带"多次出现"提示，提醒 LLM 在原文中找到该值的所有出现位置。
    """
    lines: list[str] = []
    for i, item in enumerate(items, 1):
        value = item["value"]
        value_str = json.dumps(value, ensure_ascii=False) if not isinstance(value, str) else value
        lines.append(f"{i}. 路径: {item['path']}")
        lines.append(f"   提取值: {value_str}")
        if item.get("constraint"):
            lines.append(f"   值约束: {item['constraint']}")
        if item.get("transform"):
            lines.append(f"   转换规则: {item['transform']}")
        lines.append(f"   【注意：该提取值可能在原文中出现多次，请在每次出现的位置都进行标注】")
        lines.append("")
    return "\n".join(lines)
