#!/usr/bin/env python3
"""
Batch comparison: RAG vs Wiki using vLLM on NPU (fast!)
Uses OpenAI-compatible API served by vLLM on Ascend NPU
"""
import json
import sys
import time
import urllib.request
from datetime import datetime
from pathlib import Path

sys.path.insert(0, '/home/67/knowledge/maritime')

from search_engine import search

QUESTIONS_FILE = Path('/home/67/knowledge/maritime/test_200_questions.json')
RESULTS_FILE = Path('/home/67/knowledge/maritime/batch_results_npu.json')
WIKI_DIR = '/home/67/lambdagentpaas/agentexample/qaagent67wiki/wiki'

# vLLM API on NPU
VLLM_URL = "http://127.0.0.1:8000/v1/chat/completions"

def llm_chat(messages, max_tokens=4096, temperature=0.3):
    """Call vLLM OpenAI-compatible API"""
    body = json.dumps({
        "model": "qwen2.5-32b",
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }, ensure_ascii=False).encode('utf-8')

    req = urllib.request.Request(
        VLLM_URL,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    with urllib.request.urlopen(req, timeout=300) as resp:
        data = json.loads(resp.read())
        return data["choices"][0]["message"]["content"].strip()


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

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

## 参考文档
{context}

## 问题
{question}

## 回答"""

    answer = llm_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),
    }


def wiki_answer(question, top_k=5):
    t0 = time.time()
    wiki_path = Path(WIKI_DIR)

    wiki_pages = []
    for subdir in ['sources', 'entities', 'topics', 'analyses']:
        d = wiki_path / subdir
        if d.exists():
            for f in d.glob('*.md'):
                try:
                    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({'name': f.stem, 'content': content[:1500], 'score': match_score})
                except:
                    pass

    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"""请阅读以下文档片段，提炼出与问题相关的核心知识点。
用结构化的方式组织，标注来源。

文档:
{compile_context}

问题: {question}

请输出结构化的知识摘要:"""
        compiled = llm_chat([{'role': 'user', 'content': compile_prompt}])
        wiki_context = compiled
        source_type = 'compiled'

        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 = llm_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 save_results(results, meta):
    with open(RESULTS_FILE, 'w', encoding='utf-8') as f:
        json.dump({'meta': meta, 'results': results}, f, ensure_ascii=False, indent=2)


def main():
    # Test vLLM connection
    print("Testing vLLM NPU connection...")
    try:
        test = llm_chat([{'role': 'user', 'content': 'Hello, respond with OK'}], max_tokens=10)
        print(f"vLLM OK: {test}")
    except Exception as e:
        print(f"vLLM connection failed: {e}")
        print("Make sure vLLM is running: bash /home/67/start_vllm_npu.sh")
        sys.exit(1)

    with open(QUESTIONS_FILE) as f:
        data = json.load(f)
    questions = data['questions']

    # Resume support
    existing = []
    done_ids = set()
    if RESULTS_FILE.exists():
        with open(RESULTS_FILE) as f:
            prev = json.load(f)
        existing = prev.get('results', [])
        done_ids = {r['question_id'] for r in existing}
        print(f"Resuming: {len(done_ids)} already done")

    results = list(existing)
    total = len(questions)
    remaining = [q for q in questions if q['id'] not in done_ids]

    print(f"Total: {total}, Remaining: {len(remaining)}")
    print(f"NPU inference: estimated ~3-5s per answer, total ~{len(remaining) * 8 / 60:.0f} min")
    print("=" * 60)

    start_time = time.time()

    for idx, q in enumerate(remaining):
        qnum = len(done_ids) + idx + 1
        print(f"\n[{qnum}/{total}] [{q['type']}/{q.get('difficulty','?')}] {q['question'][:60]}...")

        record = {
            'question_id': q['id'],
            'question': q['question'],
            'type': q['type'],
            'difficulty': q.get('difficulty', 'medium'),
            'category': q.get('category', ''),
        }

        # RAG
        try:
            print(f"  RAG...", end=' ', flush=True)
            rag = rag_answer(q['question'])
            record['rag'] = rag
            print(f"{rag['elapsed']}s, {len(rag['answer'])} chars")
        except Exception as e:
            record['rag'] = {'answer': f'ERROR: {e}', 'elapsed': 0, 'error': True}
            print(f"ERROR: {e}")

        # Wiki
        try:
            print(f"  Wiki...", end=' ', flush=True)
            wiki = wiki_answer(q['question'])
            record['wiki'] = wiki
            print(f"{wiki['elapsed']}s, {len(wiki['answer'])} chars, src={wiki.get('source_type','?')}")
        except Exception as e:
            record['wiki'] = {'answer': f'ERROR: {e}', 'elapsed': 0, 'error': True}
            print(f"ERROR: {e}")

        results.append(record)

        # Save every 10 questions
        if (idx + 1) % 10 == 0 or idx == len(remaining) - 1:
            elapsed_total = time.time() - start_time
            avg_per_q = elapsed_total / (idx + 1)
            eta = avg_per_q * (len(remaining) - idx - 1)
            meta = {
                'engine': 'vLLM on Ascend 910B4 NPU',
                'model': 'Qwen2.5-32B-Instruct',
                'completed': len(results),
                'total': total,
                'elapsed_minutes': round(elapsed_total / 60, 1),
                'eta_minutes': round(eta / 60, 1),
            }
            save_results(results, meta)
            print(f"  [Saved] {len(results)}/{total}, ETA: {meta['eta_minutes']}min")

    total_time = time.time() - start_time
    meta = {
        'engine': 'vLLM on Ascend 910B4 NPU',
        'model': 'Qwen2.5-32B-Instruct',
        'completed_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        'completed': len(results),
        'total': total,
        'total_minutes': round(total_time / 60, 1),
    }
    save_results(results, meta)
    print(f"\n{'=' * 60}")
    print(f"DONE! {len(results)} questions in {meta['total_minutes']} min")
    print(f"Results: {RESULTS_FILE}")


if __name__ == '__main__':
    main()