"""
提示词模板加载器
从prompt.yaml中加载提示词模板
"""

from __future__ import annotations

import csv
import re
from pathlib import Path
from typing import Dict, Any, List
from collections import OrderedDict
import yaml


class PromptLoader:
    """提示词模板加载器"""

    def __init__(self, prompt_file: Path | None = None, csv_file: Path | None = None):
        """
        初始化提示词加载器
        
        参数:
            prompt_file: 提示词文件路径，如果为None则使用默认路径
            csv_file: CSV分类标准文件路径，如果为None则使用默认路径
        """
        if prompt_file is None:
            prompt_file = Path(__file__).parent.parent / "config" / "prompt.yaml"
        self._prompt_file = Path(prompt_file)
        
        if csv_file is None:
            csv_file = Path(__file__).parent.parent / "config" / "StandardCategoryTable.csv"
        self._csv_file = Path(csv_file)
        
        self._prompts: Dict[str, Any] = {}
        self._classification_standards: str = ""
        self._standard_categories: List[str] = []  # 标准类别列表
        self._load()
        self._load_classification_standards()

    def _load(self):
        """加载提示词文件"""
        if not self._prompt_file.exists():
            raise FileNotFoundError(f"提示词文件不存在: {self._prompt_file}")
        with self._prompt_file.open("r", encoding="utf-8") as f:
            self._prompts = yaml.safe_load(f) or {}

    def _load_classification_standards(self):
        """从CSV文件加载分类标准（新格式：StandardCategoryTable.csv）"""
        if not self._csv_file.exists():
            raise FileNotFoundError(f"分类标准CSV文件不存在: {self._csv_file}")
        
        # 使用OrderedDict保持一级目录的顺序
        level1_to_level2: Dict[str, List[str]] = OrderedDict()
        
        with self._csv_file.open("r", encoding="utf-8-sig") as f:  # 使用 utf-8-sig 自动处理 BOM
            reader = csv.DictReader(f)
            for row in reader:
                # 新CSV格式：first_code, first_name, second_code, second_name, ...
                level1 = (row.get("first_name") or "").strip()
                level2 = (row.get("second_name") or "").strip()
                
                # 跳过空的一级目录
                if not level1:
                    continue
                
                # 初始化一级目录
                if level1 not in level1_to_level2:
                    level1_to_level2[level1] = []
                
                # 只添加非空且不重复的二级目录
                if level2 and level2 not in level1_to_level2[level1]:
                    level1_to_level2[level1].append(level2)
        
        # 生成分类标准文本
        # 中文数字映射
        chinese_numbers = ["一", "二", "三", "四", "五", "六", "七", "八", "九", "十"]
        
        standards_lines = []
        for idx, (level1, level2_list) in enumerate(level1_to_level2.items(), start=0):
            if idx < len(chinese_numbers):
                number_prefix = chinese_numbers[idx]
            else:
                number_prefix = str(idx + 1)
            
            # 统计二级目录数量
            level2_count = len(level2_list)
            level2_text = "、".join(level2_list)
            
            # 简化格式，只包含核心标题和二级目录列表
            if level2_count > 0:
                standards_lines.append(
                    f"    - {number_prefix}、{level1}（包含{level2_text}等{level2_count}个方面）"
                )
            else:
                standards_lines.append(
                    f"    - {number_prefix}、{level1}"
                )
        
        self._classification_standards = "\n".join(standards_lines)
        
        # 保存标准类别列表（从CSV中提取的一级目录）
        self._standard_categories = list(level1_to_level2.keys())

    def get_secondary_standards(self, first_category_name: str) -> str:
        """
        获取指定一级分类的二级分类标准

        参数:
            first_category_name: 一级分类名称（如"编制依据"）

        返回:
            二级分类标准文本（带编号列表）
        """
        if not self._csv_file.exists():
            raise FileNotFoundError(f"分类标准CSV文件不存在: {self._csv_file}")

        secondary_standards = []
        seen_names = set()

        with self._csv_file.open("r", encoding="utf-8-sig") as f:
            reader = csv.DictReader(f)
            idx = 1
            for row in reader:
                level1 = (row.get("first_name") or "").strip()
                level2_name = (row.get("second_name") or "").strip()

                # 只匹配指定的一级分类
                if level1 != first_category_name:
                    continue

                # 去重
                if level2_name and level2_name not in seen_names:
                    secondary_standards.append(f"{idx}. {level2_name}")
                    seen_names.add(level2_name)
                    idx += 1

        if not secondary_standards:
            return "（该一级分类下无二级分类标准）"

        return "\n".join(secondary_standards)

    def get_secondary_mapping(self, first_category_name: str) -> Dict[int, Dict[str, str]]:
        """
        获取指定一级分类的二级分类编号映射

        参数:
            first_category_name: 一级分类名称（如"编制依据"）

        返回:
            编号到二级分类信息的映射: {
                1: {"name": "法律法规", "code": "LawsAndRegulations"},
                2: {"name": "标准规范", "code": "StandardsAndSpecifications"},
                ...
            }
        """
        if not self._csv_file.exists():
            raise FileNotFoundError(f"分类标准CSV文件不存在: {self._csv_file}")

        mapping = {}
        seen_names = set()

        with self._csv_file.open("r", encoding="utf-8-sig") as f:
            reader = csv.DictReader(f)
            idx = 1
            for row in reader:
                level1 = (row.get("first_name") or "").strip()
                level2_name = (row.get("second_name") or "").strip()
                level2_code = (row.get("second_code") or "").strip()

                # 只匹配指定的一级分类
                if level1 != first_category_name:
                    continue

                # 去重
                if level2_name and level2_name not in seen_names:
                    mapping[idx] = {
                        "name": level2_name,
                        "code": level2_code,
                    }
                    seen_names.add(level2_name)
                    idx += 1

        return mapping

    def get_standard_categories(self) -> List[str]:
        """
        获取标准类别列表
        
        返回:
            标准类别名称列表
        """
        return self._standard_categories.copy()
    
    def get_template(self, template_name: str) -> Dict[str, str]:
        """
        获取提示词模板
        
        参数:
            template_name: 模板名称（例如 "toc_classification"）
            
        返回:
            包含system和user_template的字典
        """
        template_config = self._prompts.get(template_name, {})
        return {
            "system": template_config.get("system", ""),
            "user_template": template_config.get("user_template", "")
        }

    def render(self, template_name: str, **kwargs) -> Dict[str, str]:
        """
        渲染提示词模板
        
        参数:
            template_name: 模板名称
            **kwargs: 模板变量
            
        返回:
            包含system和user消息的字典
        """
        template = self.get_template(template_name)
        
        # 渲染user模板
        user_content = template["user_template"]
        
        # 如果模板中包含 classification_standards 占位符，自动注入分类标准
        if "{{ classification_standards }}" in user_content or "{{classification_standards}}" in user_content:
            kwargs.setdefault("classification_standards", self._classification_standards)
        
        # 替换模板变量 {{ variable }}
        def replace_var(match):
            var_name = match.group(1).strip()
            return str(kwargs.get(var_name, match.group(0)))
        
        user_content = re.sub(r"\{\{\s*(\w+)\s*\}\}", replace_var, user_content)
        
        return {
            "system": template["system"],
            "user": user_content
        }



