#!/usr/bin/env python3
"""
Wiki Full Compilation: Read domain 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

from config import cfg
PROCESSED_DIR = cfg.processed_dir
WIKI_DIR = cfg.wiki_dir
PROGRESS_FILE = WIKI_DIR / '.compile_progress.json'

# LLM API endpoints — read from env vars so the job runner can override them.
# Priority: VLLM_URL env (set by knowledge.py job runner) → Ollama default → vLLM fallback
# NOTE: never default to port 8000 — that's the AgentPaaS FastAPI server,
#       which returns index.html (SPA catch-all) with HTTP 200, fooling detect_api().
VLLM_URL  = os.environ.get('VLLM_URL',  'http://127.0.0.1:8001/v1/chat/completions')
OLLAMA_URL = os.environ.get('OLLAMA_URL', 'http://127.0.0.1:11434/v1/chat/completions')

API_URL = None  # auto-detect

def _is_llm_endpoint(url: str, timeout: int = 10) -> bool:
    """Return True only if `url` responds like an OpenAI-compatible LLM.
    We send a minimal chat request and verify the response contains 'choices'
    or a known Ollama/vLLM error (model not found etc.) — NOT HTML.
    """
    try:
        body = json.dumps({
            "model": os.environ.get('LLM_MODEL', 'qwen2.5:7b'),
            "messages": [{"role": "user", "content": "hi"}],
            "max_tokens": 3,
        }).encode('utf-8')
        req = urllib.request.Request(
            url, data=body,
            headers={"Content-Type": "application/json"}, method="POST"
        )
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read()
            # Must be valid JSON and look like an LLM response (not HTML)
            if not raw or raw[:1] in (b'<', b'!'):
                return False
            data = json.loads(raw)
            # OpenAI format: {"choices": [...]} or Ollama error: {"error": "..."}
            return "choices" in data or "error" in data or "message" in data
    except urllib.error.HTTPError as e:
        # 4xx from a real LLM (model not loaded, bad request) — server IS there
        if e.code in (400, 404, 422, 503):
            return True
        return False
    except Exception:
        return False


def detect_api():
    global API_URL
    candidates = list(dict.fromkeys([VLLM_URL, OLLAMA_URL]))  # dedup, preserve order
    for url in candidates:
        if _is_llm_endpoint(url):
            API_URL = url
            label = "Ollama" if "11434" in url or "11435" in url else "vLLM/API"
            print(f"Using: {label} ({url})")
            return
    print("ERROR: No LLM API available!")
    print(f"  Tried: {candidates}")
    print("  → Please start Ollama: brew install ollama && ollama serve && ollama pull qwen2.5:7b")
    sys.exit(1)


def llm_chat(prompt, max_tokens=2048, temperature=0.2):
    model = os.environ.get('LLM_MODEL', 'qwen2.5:7b')
    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 domain 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)

    # WIKI_REBUILD=1 → clear progress file and all generated pages for a clean rebuild
    if os.environ.get('WIKI_REBUILD') == '1':
        if PROGRESS_FILE.exists():
            PROGRESS_FILE.unlink()
            print("[Rebuild] 已清除进度文件")
        for d in ['sources', 'entities', 'topics']:
            for f in (WIKI_DIR / d).glob('*.md'):
                f.unlink()
        print("[Rebuild] 已清除旧 wiki 页面，开始全量重建")

    # Load progress
    progress = load_progress()
    # Separate truly-done from errored: only skip docs with generated source pages
    compiled_pages = {f.stem for f in (WIKI_DIR / 'sources').glob('*.md')}
    done_set = compiled_pages  # use actual output files as ground truth, not the progress file
    progress['done'] = [d for d in progress['done'] if d in compiled_pages]

    # 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}")
            # Do NOT add to done — failed docs will be retried on next incremental run
            # (only successfully compiled docs with a source page count as done)

        # 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()