#!/usr/bin/env python3
"""
海事问答智能体对比实验
RAG 模式 vs Wiki 模式
"""
import json
import time
import sys
sys.path.insert(0, '/home/67/knowledge/maritime')
sys.path.insert(0, '/home/67/lambdagentpaas')
# lambdagent migrated to src-layout — add the src dir so 'from lambdagent.X'
# still resolves when running outside an editable-install venv.
sys.path.insert(0, '/home/67/lambdagentpaas/lambdagent/src')

from search_engine import search
from lambdagent.providers import create_provider

# 初始化 LLM
provider = create_provider('ollama', model='qwen2.5:32b', timeout=600)

SEP = '=' * 70
LINE = '-' * 70

# RAG 模式
def rag_answer(question, top_k=5):
    t0 = time.time()
    results = search(question, top_k=top_k)
    context_parts = []
    for i, r in enumerate(results):
        context_parts.append(f'[doc{i+1}: {r["source"]}]\n{r["text"]}')
    context = '\n\n'.join(context_parts)

    prompt = f"""你是一个海事领域问答助手。请严格基于以下参考文档回答问题。
如果文档中没有相关信息，请说明"文档中未找到相关信息"。
回答要具体，引用文档中的原文，标注来源 [来源: 文件名]。

## 参考文档

{context}

## 问题

{question}

## 回答"""

    answer = provider.chat([{'role': 'user', 'content': prompt}])
    elapsed = time.time() - t0
    return {
        'answer': answer,
        'sources': [r['source'] for r in results],
        'elapsed': round(elapsed, 1),
        'chunks_used': len(results),
    }


# Wiki 模式
WIKI_DIR = '/home/67/lambdagentpaas/agentexample/qaagent67wiki/wiki'

def wiki_answer(question, top_k=5):
    t0 = time.time()
    from pathlib import Path

    wiki_path = Path(WIKI_DIR)

    # Step 1: 检查 wiki 中是否有相关页面
    wiki_pages = []
    for subdir in ['sources', 'entities', 'topics', 'analyses']:
        d = wiki_path / subdir
        if d.exists():
            for f in d.glob('*.md'):
                content = f.read_text(encoding='utf-8')
                q_chars = set(question)
                match_score = sum(1 for c in q_chars if c in content)
                if match_score > len(question) * 0.3:
                    wiki_pages.append({
                        'path': str(f),
                        'name': f.stem,
                        'content': content[:1500],
                        'score': match_score,
                    })

    wiki_pages.sort(key=lambda x: x['score'], reverse=True)
    wiki_pages = wiki_pages[:3]

    if wiki_pages:
        wiki_context = '\n\n'.join([
            f'[Wiki: {p["name"]}]\n{p["content"]}'
            for p in wiki_pages
        ])
        source_type = 'wiki'
    else:
        results = search(question, top_k=top_k)
        compile_context = '\n\n'.join([
            f'[doc: {r["source"]}]\n{r["text"]}'
            for r in results
        ])

        compile_prompt = f"""请阅读以下文档片段，提炼出与问题相关的核心知识点。
用结构化的方式组织，标注来源。这将作为 wiki 知识页面保存。

文档:
{compile_context}

问题: {question}

请输出结构化的知识摘要:"""

        compiled = provider.chat([{'role': 'user', 'content': compile_prompt}])
        wiki_context = compiled
        source_type = 'compiled'

        # 保存到 wiki analyses 目录
        safe_name = question[:30].replace('/', '_').replace(' ', '_')
        analysis_path = wiki_path / 'analyses' / f'{safe_name}.md'
        analysis_path.parent.mkdir(parents=True, exist_ok=True)
        with open(analysis_path, 'w', encoding='utf-8') as f:
            f.write(f'# {question}\n\n{compiled}\n')

    answer_prompt = f"""你是一个海事领域 wiki 知识库问答助手。
基于以下已编译的 wiki 知识回答问题。回答要具体，标注来源。

## Wiki 知识

{wiki_context}

## 问题

{question}

## 回答"""

    answer = provider.chat([{'role': 'user', 'content': answer_prompt}])
    elapsed = time.time() - t0
    return {
        'answer': answer,
        'source_type': source_type,
        'wiki_pages_used': len(wiki_pages),
        'elapsed': round(elapsed, 1),
    }


# 运行对比
def run_comparison():
    with open('/home/67/knowledge/maritime/test_questions.json') as f:
        data = json.load(f)

    questions = data['questions']
    results = []

    print(SEP)
    print('  Maritime QA Comparison: RAG vs Wiki')
    print(SEP)

    for q in questions:
        print(f'\n{LINE}')
        print(f'Q{q["id"]} [{q["difficulty"]}] {q["question"]}')
        print(LINE)

        # RAG
        print('\n[RAG] answering...')
        try:
            rag = rag_answer(q['question'])
            print(f'  time: {rag["elapsed"]}s | chunks: {rag["chunks_used"]}')
            print(f'  sources: {", ".join(rag["sources"][:3])}')
            ans_preview = rag["answer"][:300]
            print(f'  answer: {ans_preview}{"..." if len(rag["answer"]) > 300 else ""}')
        except Exception as e:
            rag = {'answer': f'ERROR: {e}', 'elapsed': 0}
            print(f'  ERROR: {e}')

        # Wiki
        print('\n[Wiki] answering...')
        try:
            wiki = wiki_answer(q['question'])
            print(f'  time: {wiki["elapsed"]}s | source: {wiki.get("source_type", "wiki")}')
            ans_preview = wiki["answer"][:300]
            print(f'  answer: {ans_preview}{"..." if len(wiki["answer"]) > 300 else ""}')
        except Exception as e:
            wiki = {'answer': f'ERROR: {e}', 'elapsed': 0}
            print(f'  ERROR: {e}')

        results.append({
            'question': q,
            'rag': rag,
            'wiki': wiki,
        })

    out_path = '/home/67/knowledge/maritime/comparison_results.json'
    with open(out_path, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)

    print(f'\n{SEP}')
    print(f'Done! Results saved to: {out_path}')
    print(f'\nSummary:')
    print(f'  {"Question":<40} {"RAG":>8} {"Wiki":>8}')
    for r in results:
        q_text = r['question']['question'][:36]
        print(f'  {q_text:<40} {r["rag"]["elapsed"]:>6.1f}s {r["wiki"]["elapsed"]:>6.1f}s')


if __name__ == '__main__':
    run_comparison()