#!/usr/bin/env python3
"""
Step 1: Generate 200 test questions from maritime documents
Uses document titles + sampled content to create diverse questions
"""
import json
import os
import random
import re
from pathlib import Path

PROCESSED_DIR = Path('/home/67/knowledge/maritime/processed')
OUTPUT = Path('/home/67/knowledge/maritime/test_200_questions.json')

# Collect all document files with their content
docs = []
for f in sorted(PROCESSED_DIR.glob('*.txt')):
    text = f.read_text(encoding='utf-8')
    if '[scan PDF' in text or len(text) < 100:
        continue
    docs.append({'name': f.stem, 'text': text, 'length': len(text)})

print(f"Total docs with content: {len(docs)}")

# Categorize documents by type from filename patterns
categories = {
    'crew': [],        # ship crew related
    'safety': [],      # maritime safety
    'pollution': [],   # environmental protection
    'vessel': [],      # vessel management
    'port': [],        # port management
    'law': [],         # laws and regulations
    'policy': [],      # policy documents
    'report': [],      # development reports
    'certificate': [], # certificates and licenses
    'other': [],
}

for d in docs:
    name = d['name'].lower()
    if any(k in name for k in ['船员', '海员', '适任', '培训']):
        categories['crew'].append(d)
    elif any(k in name for k in ['安全', '事故', '救助', '搜救']):
        categories['safety'].append(d)
    elif any(k in name for k in ['污染', '环保', '排放', '防治']):
        categories['pollution'].append(d)
    elif any(k in name for k in ['船舶', '船', '航行']):
        categories['vessel'].append(d)
    elif any(k in name for k in ['港口', '港', '码头', '锚地']):
        categories['port'].append(d)
    elif any(k in name for k in ['法', '条例', '规定', '办法', '规则']):
        categories['law'].append(d)
    elif any(k in name for k in ['通知', '意见', '方案', '公告']):
        categories['policy'].append(d)
    elif any(k in name for k in ['报告', '统计', '发展']):
        categories['report'].append(d)
    elif any(k in name for k in ['证书', '许可', '资质', '审批']):
        categories['certificate'].append(d)
    else:
        categories['other'].append(d)

for cat, ds in categories.items():
    print(f"  {cat}: {len(ds)} docs")

# Question templates for different types
FACT_TEMPLATES = [
    "根据{doc}，{topic}的主要内容是什么？",
    "{doc}中规定的{topic}有哪些要求？",
    "请说明{doc}中关于{topic}的具体规定。",
    "{topic}的法律依据是什么？",
    "{doc}的施行日期和适用范围是什么？",
    "{doc}中对{topic}是如何定义的？",
    "根据相关文件，{topic}需要满足哪些条件？",
    "{topic}的申请流程是怎样的？",
    "{doc}对违反{topic}的行为如何处罚？",
    "请列举{doc}中提到的{topic}的分类。",
]

REASONING_TEMPLATES = [
    "对比{doc1}和{doc2}在{topic}方面的异同。",
    "{topic}在近年来有什么变化趋势？",
    "从法律层级角度分析{topic}的规范体系。",
    "{topic}对航运企业有哪些实际影响？",
    "请综合分析{topic}存在的问题和改进建议。",
    "不同地区在{topic}方面的实施情况有何差异？",
    "{topic}与国际公约的要求有何对比？",
    "请评估{topic}的执行效果和存在的挑战。",
]

MULTI_DOC_TEMPLATES = [
    "请综合多个文件说明{topic}的完整法规框架。",
    "关于{topic}，不同层级的法规分别有哪些要求？",
    "请梳理{topic}的政策演变过程。",
    "请总结所有涉及{topic}的文件的核心要点。",
]

# Extract topics from document names
def extract_topics(doc_name):
    topics = []
    # Remove common prefixes like numbering and dates
    clean = re.sub(r'^[\d\.\-]+', '', doc_name)
    clean = re.sub(r'\d{4}年\d+月\d+日起施行', '', clean)
    clean = re.sub(r'(全国人大|国务院|交通运输部|海事局|部海事局).*?--', '', clean)
    clean = re.sub(r'--+', '', clean)
    clean = clean.strip()

    # Key maritime topics to search for
    topic_keywords = [
        '船员管理', '船舶检验', '船舶登记', '船舶安全', '海上交通',
        '港口管理', '航道管理', '海洋环保', '防污染', '船员培训',
        '适任证书', '船舶进出港', '定线制', '报告制', '引航',
        '危险货物', '防台', '搜救', '海事调查', '船舶安全检查',
        '通航安全', '水上交通', '海员', '船长', '轮机',
    ]

    for kw in topic_keywords:
        if kw in doc_name or kw in clean:
            topics.append(kw)

    if not topics and clean:
        # Use the cleaned name as topic
        if len(clean) > 5:
            topics.append(clean[:30])

    return topics

# Generate questions
questions = []
qid = 0

# 1. Fact questions from each category (120 questions)
for cat_name, cat_docs in categories.items():
    if not cat_docs:
        continue
    # Sample docs from this category
    sample_size = min(len(cat_docs), max(3, 120 // len([c for c in categories.values() if c])))
    sampled = random.sample(cat_docs, min(sample_size, len(cat_docs)))

    for doc in sampled:
        topics = extract_topics(doc['name'])
        if not topics:
            topics = [cat_name]

        topic = random.choice(topics)
        template = random.choice(FACT_TEMPLATES)

        q = template.format(doc=doc['name'][:60], topic=topic)
        qid += 1
        questions.append({
            'id': qid,
            'type': 'fact',
            'difficulty': random.choice(['easy', 'medium']),
            'category': cat_name,
            'question': q,
            'source_doc': doc['name'],
        })

# 2. Reasoning questions (50 questions)
all_topics = ['船员管理', '船舶安全', '海洋环境保护', '港口管理', '航道安全',
              '危险货物运输', '船舶检验', '海事调查', '引航管理', '船员培训',
              '船舶登记', '海上搜救', '通航安全', '防台风', '船舶防污染',
              '水上交通事故', '船员适任', '船舶技术', '航运市场', '海事执法']

for i in range(50):
    topic = random.choice(all_topics)

    if random.random() < 0.4 and len(docs) > 1:
        # Cross-document comparison
        d1, d2 = random.sample(docs[:200], 2)
        template = random.choice(REASONING_TEMPLATES[:2])
        q = template.format(doc1=d1['name'][:40], doc2=d2['name'][:40], topic=topic)
    else:
        template = random.choice(REASONING_TEMPLATES)
        q = template.format(topic=topic, doc1='', doc2='')
        q = q.replace('和在', '在').replace('  ', ' ')

    qid += 1
    questions.append({
        'id': qid,
        'type': 'reasoning',
        'difficulty': random.choice(['medium', 'hard']),
        'category': 'cross-doc',
        'question': q,
    })

# 3. Multi-doc synthesis questions (30 questions)
synthesis_topics = ['船员权益保障', '船舶安全管理体系', '海事行政处罚',
                    '船舶污染防治', '港口危险货物管理', '海上应急救援',
                    '船员职业发展', '内河航运管理', '国际航行船舶管理',
                    '海事信息化建设', '船舶技术标准', '海上通信安全',
                    '船舶保险制度', '航运企业安全管理', '海事法律体系']

for i in range(30):
    topic = random.choice(synthesis_topics)
    template = random.choice(MULTI_DOC_TEMPLATES)
    q = template.format(topic=topic)
    qid += 1
    questions.append({
        'id': qid,
        'type': 'synthesis',
        'difficulty': 'hard',
        'category': 'synthesis',
        'question': q,
    })

# Shuffle and trim to 200
random.shuffle(questions)
questions = questions[:200]
# Re-number
for i, q in enumerate(questions):
    q['id'] = i + 1

# Save
with open(OUTPUT, 'w', encoding='utf-8') as f:
    json.dump({
        'description': 'Maritime QA benchmark - 200 questions',
        'generated_at': __import__('datetime').datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        'stats': {
            'total': len(questions),
            'fact': sum(1 for q in questions if q['type'] == 'fact'),
            'reasoning': sum(1 for q in questions if q['type'] == 'reasoning'),
            'synthesis': sum(1 for q in questions if q['type'] == 'synthesis'),
        },
        'questions': questions,
    }, f, ensure_ascii=False, indent=2)

print(f"\nGenerated {len(questions)} questions:")
print(f"  Fact: {sum(1 for q in questions if q['type'] == 'fact')}")
print(f"  Reasoning: {sum(1 for q in questions if q['type'] == 'reasoning')}")
print(f"  Synthesis: {sum(1 for q in questions if q['type'] == 'synthesis')}")
print(f"\nSaved to: {OUTPUT}")
print(f"\nSample questions:")
for q in questions[:5]:
    print(f"  [{q['type']}/{q['difficulty']}] {q['question']}")