#!/usr/bin/env python3
"""
在知识库中搜索关键词（支持标准分类和自定义分类）

用法:
    uv run search_info.py <关键词>
    uv run search_info.py -i <关键词>    # 忽略大小写

示例:
    uv run search_info.py 玄天剑
    uv run search_info.py 技能
    uv run search_info.py -i sword
"""

import re
import sys
from pathlib import Path


# 标准分类
STANDARD_TYPES = ['角色', '物品', '场景', '事件']


def parse_info_structure():
    """
    解析 现有信息/信息结构.md 获取自定义分类信息

    返回: dict, 格式为 {分类名: {'type': 'directory'|'single_file', 'path': 路径}}
    """
    structure_file = Path.cwd() / "现有信息" / "信息结构.md"
    custom_categories = {}

    if not structure_file.exists():
        return custom_categories

    try:
        content = structure_file.read_text(encoding='utf-8')

        # 查找 "## 自定义分类" 部分
        custom_section_match = re.search(r'## 自定义分类\s*\n(.*?)(?=\n## |\Z)', content, re.DOTALL)
        if not custom_section_match:
            return custom_categories

        custom_section = custom_section_match.group(1)

        # 如果自定义分类部分只有"无"，则返回空
        if custom_section.strip() == '无' or '无自定义分类' in custom_section:
            return custom_categories

        # 查找所有 ### 开头的分类定义
        category_pattern = r'### ([^\n]+)\s*\n(.*?)(?=\n### |\Z)'
        categories = re.findall(category_pattern, custom_section, re.DOTALL)

        for name, details in categories:
            name = name.strip()
            if not name or name.startswith('['):  # 跳过模板占位符
                continue

            # 解析路径和结构类型
            path_match = re.search(r'\*\*路径\*\*:\s*`([^`]+)`', details)
            type_match = re.search(r'\*\*结构类型\*\*:\s*(目录|单文件)', details)

            if path_match:
                path = path_match.group(1)
                structure_type = 'single_file' if (type_match and type_match.group(1) == '单文件') else 'directory'

                # 确定实际路径
                if path.endswith('/'):
                    structure_type = 'directory'
                elif path.endswith('.md'):
                    structure_type = 'single_file'

                custom_categories[name] = {
                    'type': structure_type,
                    'path': path
                }

    except Exception:
        pass

    return custom_categories


def search_in_file(file_path: Path, keyword: str, ignore_case: bool = False) -> list[str]:
    """
    在文件中搜索关键词，返回匹配的行及其上下文

    返回: 匹配的文本片段列表
    """
    try:
        content = file_path.read_text(encoding='utf-8')
        lines = content.split('\n')

        matches = []
        search_keyword = keyword.lower() if ignore_case else keyword

        for i, line in enumerate(lines):
            line_to_search = line.lower() if ignore_case else line
            if search_keyword in line_to_search:
                # 获取上下文（前后各1行）
                start = max(0, i - 1)
                end = min(len(lines), i + 2)
                context_lines = lines[start:end]

                # 高亮关键词
                highlighted = []
                for ctx_line in context_lines:
                    if ignore_case:
                        # 忽略大小写时的高亮
                        pattern = re.compile(re.escape(keyword), re.IGNORECASE)
                        highlighted_line = pattern.sub(f">>>{keyword}<<<", ctx_line)
                    else:
                        highlighted_line = ctx_line.replace(keyword, f">>>{keyword}<<<")
                    highlighted.append(highlighted_line)

                context = '\n'.join(highlighted)
                matches.append(context)

        return matches

    except Exception:
        return []


def get_category_type(info_type: str, base_path: Path, custom_categories: dict) -> str:
    """获取分类的类型"""
    if info_type in STANDARD_TYPES:
        return 'directory'

    if info_type in custom_categories:
        return custom_categories[info_type]['type']

    # 尝试自动检测
    dir_path = base_path / info_type
    if dir_path.is_dir():
        return 'directory'

    file_path = base_path / f"{info_type}.md"
    if file_path.is_file():
        return 'single_file'

    return 'unknown'


def search_knowledge_base(keyword: str, ignore_case: bool = False):
    """
    在整个知识库中搜索关键词（包括自定义分类）

    Args:
        keyword: 要搜索的关键词
        ignore_case: 是否忽略大小写
    """
    base_path = Path.cwd() / "现有信息"

    if not base_path.exists():
        print(f"错误: 未找到目录 '{base_path}'")
        print("请确保当前工作目录包含 '现有信息/' 文件夹")
        return

    case_note = "（忽略大小写）" if ignore_case else ""
    print(f"搜索关键词: {keyword} {case_note}")
    print()

    # 获取自定义分类
    custom_categories = parse_info_structure()

    # 构建所有要搜索的分类
    all_types = STANDARD_TYPES.copy()
    for name in custom_categories.keys():
        if name not in all_types:
            all_types.append(name)

    # 同时搜索现有信息目录下的其他目录（未在信息结构.md中定义的）
    if base_path.exists():
        for item in base_path.iterdir():
            if item.is_dir() and item.name not in all_types:
                all_types.append(item.name)

    total_matches = 0
    results = []

    for info_type in all_types:
        cat_type = get_category_type(info_type, base_path, custom_categories)

        if cat_type == 'single_file':
            # 单文件结构
            file_path = base_path / f"{info_type}.md"
            if file_path.exists():
                matches = search_in_file(file_path, keyword, ignore_case)
                if matches:
                    total_matches += len(matches)
                    results.append({
                        'type': info_type,
                        'name': info_type,
                        'path': f"现有信息/{info_type}.md",
                        'matches': matches,
                        'is_single_file': True
                    })
        else:
            # 目录结构
            type_path = base_path / info_type
            if not type_path.exists():
                continue

            md_files = sorted(type_path.glob("*.md"))

            for md_file in md_files:
                matches = search_in_file(md_file, keyword, ignore_case)

                if matches:
                    total_matches += len(matches)
                    results.append({
                        'type': info_type,
                        'name': md_file.stem,
                        'path': f"现有信息/{info_type}/{md_file.name}",
                        'matches': matches,
                        'is_single_file': False
                    })

    # 输出结果
    if results:
        for result in results:
            if result['is_single_file']:
                print(f"[{result['type']}] (单文件)")
            else:
                print(f"[{result['type']}] {result['name']}")
            print(f"文件: {result['path']}")
            print()

            # 显示匹配的内容（限制每个文件最多显示3个匹配）
            for i, match in enumerate(result['matches'][:3]):
                print(f"  匹配 {i+1}:")
                for line in match.split('\n'):
                    print(f"    {line}")
                print()

            if len(result['matches']) > 3:
                print(f"  ... 还有 {len(result['matches']) - 3} 个匹配")
                print()

            print("-" * 80)
            print()

    if total_matches == 0:
        print(f"未找到包含 '{keyword}' 的信息")
    else:
        print(f"总计找到 {total_matches} 个匹配结果")


def main():
    if len(sys.argv) < 2:
        print("用法: uv run search_info.py [-i] <关键词>")
        print()
        print("选项:")
        print("  -i: 忽略大小写")
        print()
        print("说明:")
        print("  搜索范围包括标准分类（角色/物品/场景/事件）和自定义分类")
        print()
        print("示例:")
        print("  uv run search_info.py 玄天剑")
        print("  uv run search_info.py 技能        # 会搜索自定义分类中的技能相关内容")
        print("  uv run search_info.py -i sword")
        sys.exit(1)

    # 解析参数
    ignore_case = '-i' in sys.argv
    args = [arg for arg in sys.argv[1:] if arg != '-i']

    if not args:
        print("错误: 请提供搜索关键词")
        sys.exit(1)

    keyword = args[0]
    search_knowledge_base(keyword, ignore_case)


if __name__ == '__main__':
    main()
