"""
手动反思工具

针对已完成或正在进行的 trace，手动触发反思并提供建议。

用法：
    python manual_reflect.py --trace <trace_id> --sequence <seq>
    python manual_reflect.py --trace <trace_id> --sequence <seq> --feedback "建议内容"
"""

import asyncio
import argparse
import sys
from pathlib import Path

# 添加项目根目录到 Python 路径
sys.path.insert(0, str(Path(__file__).parent))

from dotenv import load_dotenv
load_dotenv()

from agent.core.runner import AgentRunner, RunConfig
from agent.trace import FileSystemTraceStore
from agent.llm import create_openrouter_llm_call
from agent.tools.builtin.knowledge import KnowledgeConfig


def read_multiline() -> str:
    """读取多行输入，以连续两次回车（空行）结束"""
    print("\n请输入你的建议（连续输入两次回车结束）：")
    lines = []
    blank_count = 0
    while True:
        line = input()
        if line == "":
            blank_count += 1
            if blank_count >= 2:
                break
            lines.append("")
        else:
            blank_count = 0
            lines.append(line)

    # 移除末尾的空行
    while lines and lines[-1] == "":
        lines.pop()
    return "\n".join(lines)


async def main():
    parser = argparse.ArgumentParser(
        description="手动反思工具：基于用户建议对 trace 历史进行反思"
    )
    parser.add_argument("--trace", required=True, help="Trace ID")
    parser.add_argument("--sequence", type=int, required=True, help="反思截止点（包含此 sequence）")
    parser.add_argument("--feedback", help="用户建议（可选，不指定则交互式输入）")
    parser.add_argument("--store-path", default=".trace", help="TraceStore 路径")
    parser.add_argument("--model", default="claude-sonnet-4.5", help="使用的模型")
    parser.add_argument("--config", help="配置文件路径（如 examples/research/config.py），用于加载 KnowledgeConfig")
    args = parser.parse_args()

    trace_id = args.trace
    sequence = args.sequence

    # 初始化 TraceStore
    store = FileSystemTraceStore(base_path=args.store_path)

    # 验证 trace 存在
    trace = await store.get_trace(trace_id)
    if not trace:
        print(f"❌ 错误：Trace 不存在: {trace_id}")
        sys.exit(1)

    print(f"\n=== 手动反思 ===")
    print(f"Trace ID: {trace_id[:8]}...")
    print(f"状态: {trace.status}")
    print(f"总消息数: {trace.total_messages}")
    print(f"反思截止点: sequence {sequence}")

    # 获取用户建议
    if args.feedback:
        feedback = args.feedback
    else:
        feedback = read_multiline()

    if not feedback.strip():
        print("❌ 错误：未提供建议")
        sys.exit(1)

    print(f"\n建议长度: {len(feedback)} 字符")

    # 获取历史消息（从 sequence 回溯到 root）
    print(f"\n正在加载历史消息...")
    messages = await store.get_main_path_messages(trace_id, sequence)
    print(f"  加载了 {len(messages)} 条消息")

    # 转换为 OpenAI 格式（使用 to_llm_dict() 正确处理 assistant 消息的 dict 格式）
    history = []
    for msg in messages:
        if msg.role in ("user", "assistant"):
            history.append(msg.to_llm_dict())

    print(f"  转换后历史消息数: {len(history)}")

    # 构造反思 prompt（插入用户建议）
    from agent.core.prompts.knowledge import MANUAL_REFLECT_PROMPT
    reflect_prompt = MANUAL_REFLECT_PROMPT.format(user_feedback=feedback)

    # 初始化 Runner
    runner = AgentRunner(
        trace_store=store,
        llm_call=create_openrouter_llm_call(model=f"anthropic/{args.model}"),
        debug=False
    )

    # 加载 KnowledgeConfig
    knowledge_config = KnowledgeConfig()

    if args.config:
        # 从配置文件加载
        config_path = Path(args.config)
        if not config_path.exists():
            print(f"❌ 错误：配置文件不存在: {args.config}")
            sys.exit(1)

        print(f"\n正在加载配置: {args.config}")

        # 动态导入配置模块
        import importlib.util
        spec = importlib.util.spec_from_file_location("config", config_path)
        config_module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(config_module)

        if hasattr(config_module, 'RUN_CONFIG') and hasattr(config_module.RUN_CONFIG, 'knowledge'):
            knowledge_config = config_module.RUN_CONFIG.knowledge
            print(f"  owner: {knowledge_config.owner or '(默认)'}")
            print(f"  default_tags: {knowledge_config.default_tags}")
            print(f"  default_scopes: {knowledge_config.default_scopes}")
        else:
            print(f"  ⚠️  配置文件中未找到 RUN_CONFIG.knowledge，使用默认配置")

    # 使用加载的 KnowledgeConfig
    config = RunConfig(
        model=args.model,
        knowledge=knowledge_config
    )

    # 执行反思
    print(f"\n正在执行反思...")
    print(f"  模型: {args.model}")
    print(f"  Prompt 长度: {len(reflect_prompt)} 字符")

    await runner._run_reflect(
        trace_id=trace_id,
        history=history,
        config=config,
        reflect_prompt=reflect_prompt,
        source_name="manual_reflection"
    )

    print(f"\n✅ 反思完成")
    print(f"\n提示：经验已保存到知识库，可通过 knowledge_search 工具查询")


if __name__ == "__main__":
    asyncio.run(main())