#!/usr/bin/env python3
"""
Wiki Full Compilation: Read 1326 maritime docs → LLM understand → Write wiki pages
Uses vLLM on NPU for fast inference

For each document:
  1. Read extracted text
  2. LLM summarizes → write sources/<filename>.md
  3. LLM extracts entities → create/update entities/<entity>.md
  4. LLM identifies topics → create/update topics/<topic>.md
  5. Update index.md and log.md

Supports resume from interruption.
"""
import json
import os
import re
import sys
import time
import urllib.request
from datetime import datetime
from pathlib import Path

PROCESSED_DIR = Path('/data/knowledge/maritime/processed')
WIKI_DIR = Path('/app/agentexample/qaagent67wiki/wiki')
PROGRESS_FILE = WIKI_DIR / '.compile_progress.json'

# vLLM API
VLLM_URL = "http://127.0.0.1:8000/v1/chat/completions"
# Fallback to Ollama if vLLM not available
OLLAMA_URL = "http://127.0.0.1:8000/v1/chat/completions"

API_URL = None  # auto-detect

def detect_api():
    global API_URL
    for url in [VLLM_URL, OLLAMA_URL]:
        try:
            test_body = json.dumps({
                "model": "qwen2.5-32b" if "8000" in url else "qwen2.5-32b",
                "messages": [{"role": "user", "content": "hi"}],
                "max_tokens": 5,
            }).encode('utf-8')
            req = urllib.request.Request(url, data=test_body,
                                        headers={"Content-Type": "application/json"}, method="POST")
            with urllib.request.urlopen(req, timeout=30) as resp:
                API_URL = url
                engine = "vLLM NPU" if "8000" in url else "Ollama CPU"
                print(f"Using: {engine} ({url})")
                return
        except:
            continue
    print("ERROR: No LLM API available!")
    sys.exit(1)


def llm_chat(prompt, max_tokens=2048, temperature=0.2):
    model = "qwen2.5-32b" if "8000" in API_URL else "qwen2.5-32b"
    body = json.dumps({
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": temperature,
        "max_tokens": max_tokens,
    }, ensure_ascii=False).encode('utf-8')

    req = urllib.request.Request(API_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 safe_filename(name, max_len=80):
    """Convert to safe filename"""
    name = re.sub(r'[\\/:*?"<>|]', '_', name)
    name = re.sub(r'\s+', '_', name)
    return name[:max_len]


def load_progress():
    if PROGRESS_FILE.exists():
        with open(PROGRESS_FILE) as f:
            return json.load(f)
    return {'done': [], 'entities': {}, 'topics': {}, 'relations': []}


def save_progress(progress):
    with open(PROGRESS_FILE, 'w', encoding='utf-8') as f:
        json.dump(progress, f, ensure_ascii=False, indent=2)


def _extract_date_from_name(doc_name):
    """从文件名中提取日期信息"""
    # 匹配 "2017年8月31日起施行" 或 "2022年4月1日"
    m = re.search(r'(\d{4})年(\d{1,2})月(\d{1,2})日', doc_name)
    if m:
        return f"{m.group(1)}年{m.group(2)}月{m.group(3)}日"
    # 匹配 "2017年"
    m = re.search(r'(\d{4})年', doc_name)
    if m:
        return f"{m.group(1)}年"
    return "未知"


def compile_source(doc_name, text, progress):
    """Step 1: Create source summary page"""
    # Truncate long documents
    text_sample = text[:6000] if len(text) > 6000 else text

    # Extract date from filename for temporal normalization
    doc_date = _extract_date_from_name(doc_name)

    prompt = f"""请阅读以下海事领域文档，输出结构化摘要。

文档名: {doc_name}
文档日期: {doc_date}
文档内容:
{text_sample}

请按以下格式输出（不要输出其他内容）:

## 摘要
[2-4句话概括文档核心内容，必须包含具体的施行日期和文号]

## 关键信息
- [要点1]
- [要点2]
- [要点3]
（最多8个要点，每个要点尽量包含具体条款号、数字、日期）

## 时间线
[列出文档中提到的所有时间节点，格式: 绝对日期|事件]
{doc_date}|本文件施行
（**重要**: 文档中的所有相对时间必须转换为绝对时间:
  - "自发布之日起施行" → "{doc_date}|施行"
  - "近年来" → 以{doc_date}为基准，推算为具体年份范围
  - "去年" → 以{doc_date}为基准，推算具体年份
  - "本条例修订前" → 写明具体的修订前版本年份
  - "三年内" → 从{doc_date}起算，写明截止的具体年份
  禁止使用任何相对时间表述，全部转为绝对年月日）

## 实体列表
[列出文档中涉及的重要实体，每行一个，格式: 类型|名称|说明]
法规|XXX法|（含施行日期和文号）
机构|XXX局|...
概念|XXX制度|...

## 实体关系
[列出文档中实体之间的关系，每行一个，格式: 主体|关系|客体|说明]
交通运输部|制定|船员管理规定|交通运输部负责制定船员管理相关规定
海事局|隶属|交通运输部|海事局是交通运输部下属机构
船员适任证书|依据|STCW公约|船员适任证书制度依据STCW国际公约
（列出文档中所有重要的实体间关系，包括：上下级/隶属、制定/发布、依据/引用、适用/管辖、修订/废止、时间先后等）

## 相关主题
[列出文档涉及的主题领域，每行一个]
船员管理
船舶安全
...

## 元数据
- 发布机构: [如有]
- 施行日期: {doc_date}
- 文号: [如有]
- 法律层级: [全国人大/国务院/部委/海事局/其他]
- 是否现行有效: [是/否/已被修订]
- 修订历史: [如有，列出历次修订的绝对日期]"""

    result = llm_chat(prompt)

    # Write source page
    source_path = WIKI_DIR / 'sources' / f'{safe_filename(doc_name)}.md'
    source_path.parent.mkdir(parents=True, exist_ok=True)
    with open(source_path, 'w', encoding='utf-8') as f:
        f.write(f'# {doc_name}\n\n')
        f.write(f'[来源: {doc_name}]\n\n')
        f.write(result)
        f.write('\n')

    return result


def extract_entities(llm_result, doc_name, progress):
    """Step 2: Extract and update entity pages"""
    entities_found = []
    lines = llm_result.split('\n')
    in_entity_section = False

    for line in lines:
        if '## 实体列表' in line:
            in_entity_section = True
            continue
        if line.startswith('## ') and in_entity_section:
            in_entity_section = False
            continue
        if in_entity_section and '|' in line:
            parts = [p.strip() for p in line.strip('- ').split('|')]
            if len(parts) >= 2:
                etype = parts[0]
                ename = parts[1]
                edesc = parts[2] if len(parts) > 2 else ''
                if ename and len(ename) > 1:
                    entities_found.append((etype, ename, edesc))

    # Update entity pages
    for etype, ename, edesc in entities_found[:10]:  # Max 10 per doc
        entity_file = safe_filename(ename)
        entity_path = WIKI_DIR / 'entities' / f'{entity_file}.md'
        entity_path.parent.mkdir(parents=True, exist_ok=True)

        if entity_path.exists():
            # Append reference
            existing = entity_path.read_text(encoding='utf-8')
            if doc_name not in existing:
                with open(entity_path, 'a', encoding='utf-8') as f:
                    f.write(f'\n- {edesc} [来源: {doc_name}]\n')
        else:
            # Create new entity page
            with open(entity_path, 'w', encoding='utf-8') as f:
                f.write(f'# {ename}\n\n')
                f.write(f'**类型**: {etype}\n\n')
                f.write(f'## 描述\n\n{edesc}\n\n')
                f.write(f'## 相关文档\n\n')
                f.write(f'- {doc_name}\n')

        # Track in progress
        if ename not in progress['entities']:
            progress['entities'][ename] = {'type': etype, 'refs': []}
        if doc_name not in progress['entities'][ename]['refs']:
            progress['entities'][ename]['refs'].append(doc_name)

    return entities_found


def extract_topics(llm_result, doc_name, progress):
    """Step 3: Extract and update topic pages"""
    topics_found = []
    lines = llm_result.split('\n')
    in_topic_section = False

    for line in lines:
        if '## 相关主题' in line:
            in_topic_section = True
            continue
        if line.startswith('## ') and in_topic_section:
            in_topic_section = False
            continue
        if in_topic_section:
            topic = line.strip('- ').strip()
            if topic and len(topic) > 1 and not topic.startswith('#'):
                topics_found.append(topic)

    # Update topic pages
    for topic in topics_found[:5]:  # Max 5 per doc
        topic_file = safe_filename(topic)
        topic_path = WIKI_DIR / 'topics' / f'{topic_file}.md'
        topic_path.parent.mkdir(parents=True, exist_ok=True)

        if topic_path.exists():
            existing = topic_path.read_text(encoding='utf-8')
            if doc_name not in existing:
                with open(topic_path, 'a', encoding='utf-8') as f:
                    f.write(f'- [[{safe_filename(doc_name)}]] - {doc_name}\n')
        else:
            with open(topic_path, 'w', encoding='utf-8') as f:
                f.write(f'# {topic}\n\n')
                f.write(f'## 相关文档\n\n')
                f.write(f'- [[{safe_filename(doc_name)}]] - {doc_name}\n')

        if topic not in progress['topics']:
            progress['topics'][topic] = []
        if doc_name not in progress['topics'][topic]:
            progress['topics'][topic].append(doc_name)

    return topics_found


def extract_relations(llm_result, doc_name, progress):
    """Step 4: Extract entity relations and build knowledge graph"""
    relations_found = []
    lines = llm_result.split('\n')
    in_relation_section = False

    for line in lines:
        if '## 实体关系' in line:
            in_relation_section = True
            continue
        if line.startswith('## ') and in_relation_section:
            in_relation_section = False
            continue
        if in_relation_section and '|' in line:
            parts = [p.strip() for p in line.strip('- （').split('|')]
            if len(parts) >= 3:
                subject = parts[0]
                relation = parts[1]
                obj = parts[2]
                desc = parts[3] if len(parts) > 3 else ''
                if subject and relation and obj and len(subject) > 1:
                    relations_found.append({
                        'subject': subject,
                        'relation': relation,
                        'object': obj,
                        'description': desc,
                        'source': doc_name,
                    })

    # Store relations in progress
    if 'relations' not in progress:
        progress['relations'] = []
    for rel in relations_found[:15]:  # Max 15 per doc
        # Deduplicate
        existing = any(
            r['subject'] == rel['subject'] and
            r['relation'] == rel['relation'] and
            r['object'] == rel['object']
            for r in progress['relations']
        )
        if not existing:
            progress['relations'].append(rel)

    # Update entity pages with relations
    for rel in relations_found[:15]:
        for ename in [rel['subject'], rel['object']]:
            entity_file = safe_filename(ename)
            entity_path = WIKI_DIR / 'entities' / f'{entity_file}.md'
            if entity_path.exists():
                existing = entity_path.read_text(encoding='utf-8')
                if '## 关系' not in existing:
                    with open(entity_path, 'a', encoding='utf-8') as f:
                        f.write(f'\n## 关系\n\n')
                rel_line = f'- {rel["subject"]} **{rel["relation"]}** {rel["object"]}'
                if rel['description']:
                    rel_line += f' — {rel["description"]}'
                rel_line += f' [来源: {doc_name}]\n'
                existing = entity_path.read_text(encoding='utf-8')
                if rel_line.strip() not in existing:
                    with open(entity_path, 'a', encoding='utf-8') as f:
                        f.write(rel_line)

    # Write/update relations graph file
    relations_path = WIKI_DIR / 'relations.json'
    with open(relations_path, 'w', encoding='utf-8') as f:
        json.dump(progress['relations'], f, ensure_ascii=False, indent=2)

    return relations_found


def update_index(progress):
    """Update wiki index.md"""
    sources = sorted(WIKI_DIR.glob('sources/*.md'))
    entities = sorted(WIKI_DIR.glob('entities/*.md'))
    topics = sorted(WIKI_DIR.glob('topics/*.md'))
    analyses = sorted(WIKI_DIR.glob('analyses/*.md'))

    index = f"""# Maritime Wiki Index

> Auto-compiled from 1326 maritime documents by wikiagent
> Last updated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

## Sources ({len(sources)} files)

"""
    for s in sources[:50]:  # Show first 50
        index += f'- [[{s.stem}]]\n'
    if len(sources) > 50:
        index += f'\n... and {len(sources) - 50} more\n'

    index += f'\n## Entities ({len(entities)})\n\n'
    # Group by type
    entity_types = {}
    for ename, einfo in progress.get('entities', {}).items():
        etype = einfo.get('type', 'other')
        if etype not in entity_types:
            entity_types[etype] = []
        entity_types[etype].append((ename, len(einfo.get('refs', []))))

    for etype, elist in sorted(entity_types.items()):
        index += f'### {etype} ({len(elist)})\n'
        for ename, ref_count in sorted(elist, key=lambda x: -x[1])[:20]:
            index += f'- [[{safe_filename(ename)}]] ({ref_count} refs)\n'
        if len(elist) > 20:
            index += f'- ... and {len(elist) - 20} more\n'
        index += '\n'

    index += f'## Topics ({len(topics)})\n\n'
    for tname, trefs in sorted(progress.get('topics', {}).items(), key=lambda x: -len(x[1]))[:30]:
        index += f'- [[{safe_filename(tname)}]] ({len(trefs)} docs)\n'
    if len(topics) > 30:
        index += f'\n... and {len(topics) - 30} more\n'

    index += f'\n## Analyses ({len(analyses)})\n\n'
    for a in analyses[:20]:
        index += f'- [[{a.stem}]]\n'

    # Relations summary
    relations = progress.get('relations', [])
    if relations:
        index += f'\n## Relations ({len(relations)})\n\n'
        # Group by relation type
        rel_types = {}
        for r in relations:
            rt = r['relation']
            if rt not in rel_types:
                rel_types[rt] = []
            rel_types[rt].append(r)
        for rt, rels in sorted(rel_types.items(), key=lambda x: -len(x[1]))[:15]:
            index += f'### {rt} ({len(rels)})\n'
            for r in rels[:5]:
                index += f'- {r["subject"]} → {r["object"]}\n'
            if len(rels) > 5:
                index += f'- ... and {len(rels) - 5} more\n'
            index += '\n'

    index += f'\n---\nStats: {len(sources)} sources, {len(entities)} entities, {len(topics)} topics, {len(relations)} relations, {len(analyses)} analyses\n'

    with open(WIKI_DIR / 'index.md', 'w', encoding='utf-8') as f:
        f.write(index)


def update_log(doc_name, entities_count, topics_count, relations_count=0):
    """Append to log.md"""
    log_path = WIKI_DIR / 'log.md'
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M')
    entry = f'## [{timestamp}] Ingest | {doc_name}\n- Entities: {entities_count}, Topics: {topics_count}, Relations: {relations_count}\n\n'

    if log_path.exists():
        existing = log_path.read_text(encoding='utf-8')
        # Insert after header
        if '_no records_' in existing or '_no operations_' in existing:
            existing = '# Operation Log\n\n'
        with open(log_path, 'w', encoding='utf-8') as f:
            f.write(existing + entry)
    else:
        with open(log_path, 'w', encoding='utf-8') as f:
            f.write('# Operation Log\n\n' + entry)


def main():
    detect_api()

    # Ensure wiki dirs exist
    for d in ['sources', 'entities', 'topics', 'analyses']:
        (WIKI_DIR / d).mkdir(parents=True, exist_ok=True)

    # Load progress
    progress = load_progress()
    done_set = set(progress['done'])

    # Get all text files
    all_docs = sorted(PROCESSED_DIR.glob('*.txt'))
    remaining = [d for d in all_docs if d.stem not in done_set]

    total = len(all_docs)
    print(f"Total documents: {total}")
    print(f"Already compiled: {len(done_set)}")
    print(f"Remaining: {len(remaining)}")
    print("=" * 60)

    start_time = time.time()
    errors = 0

    for idx, doc_path in enumerate(remaining):
        doc_name = doc_path.stem
        qnum = len(done_set) + idx + 1

        text = doc_path.read_text(encoding='utf-8')
        if len(text) < 50:
            print(f"[{qnum}/{total}] SKIP (too short): {doc_name[:60]}")
            progress['done'].append(doc_name)
            continue

        print(f"\n[{qnum}/{total}] {doc_name[:70]}...")

        try:
            # Step 1: Compile source summary
            t0 = time.time()
            llm_result = compile_source(doc_name, text, progress)
            t1 = time.time()
            print(f"  Source page: {t1-t0:.1f}s")

            # Step 2: Extract entities
            entities = extract_entities(llm_result, doc_name, progress)
            print(f"  Entities: {len(entities)} found")

            # Step 3: Extract topics
            topics = extract_topics(llm_result, doc_name, progress)
            print(f"  Topics: {len(topics)} found")

            # Step 4: Extract relations
            relations = extract_relations(llm_result, doc_name, progress)
            print(f"  Relations: {len(relations)} found")

            # Step 5: Update log
            update_log(doc_name, len(entities), len(topics), len(relations))

            progress['done'].append(doc_name)

        except Exception as e:
            errors += 1
            print(f"  ERROR: {e}")
            # Still mark as done to avoid infinite retry
            progress['done'].append(doc_name)

        # Save progress and update index every 20 docs
        if (idx + 1) % 20 == 0 or idx == len(remaining) - 1:
            save_progress(progress)
            update_index(progress)

            elapsed = time.time() - start_time
            avg = elapsed / (idx + 1)
            eta = avg * (len(remaining) - idx - 1)
            print(f"\n  [Progress] {qnum}/{total} done | "
                  f"Elapsed: {elapsed/60:.1f}min | ETA: {eta/60:.1f}min | "
                  f"Errors: {errors} | "
                  f"Entities: {len(progress['entities'])} | "
                  f"Topics: {len(progress['topics'])} | "
                  f"Relations: {len(progress.get('relations', []))}")

    # Final update
    save_progress(progress)
    update_index(progress)

    total_time = time.time() - start_time
    sources = len(list(WIKI_DIR.glob('sources/*.md')))
    entities = len(list(WIKI_DIR.glob('entities/*.md')))
    topics = len(list(WIKI_DIR.glob('topics/*.md')))

    print(f"\n{'=' * 60}")
    print(f"Wiki Compilation Complete!")
    print(f"  Time: {total_time/60:.1f} minutes")
    print(f"  Sources: {sources}")
    print(f"  Entities: {entities}")
    print(f"  Topics: {topics}")
    print(f"  Errors: {errors}")
    print(f"  Wiki dir: {WIKI_DIR}")


if __name__ == '__main__':
    main()