#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import os
import re
from typing import Dict, List, Tuple


NODE_RE = re.compile(r'^\s*@node\s+(?P<id>[A-Za-z0-9_]+)\s*$')
TITLE_RE = re.compile(r'^\s*@title\s+"(?P<title>.*)"\s*$')
CHOICE_RE = re.compile(r'^\s*choice_\d+\s*:\s*"(?P<label>.*?)"\s*->\s*(?P<target>[A-Za-z0-9_]+)')

ENDING_ID_PREFIXES = (
    'ending_',
)
ENDING_TITLE_KEYWORDS = (
    '结局', '最终', '终极', '拯救', '守护', '和谐', '文明', '宽恕', '宇宙', '完美'
)


def parse_story(path: str) -> Tuple[Dict[str, str], Dict[str, List[Tuple[str, str]]]]:
    with open(path, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    nodes: Dict[str, str] = {}
    edges: Dict[str, List[Tuple[str, str]]] = {}

    current_id: str = ''
    i = 0
    n = len(lines)
    while i < n:
        line = lines[i].rstrip('\n')
        m_node = NODE_RE.match(line)
        if m_node:
            current_id = m_node.group('id')
            if current_id not in nodes:
                nodes[current_id] = current_id
            if current_id not in edges:
                edges[current_id] = []

            j = i + 1
            while j < n:
                l2 = lines[j].rstrip('\n')
                if NODE_RE.match(l2):
                    break
                m_title = TITLE_RE.match(l2)
                if m_title:
                    title = m_title.group('title').strip()
                    if title:
                        nodes[current_id] = title
                m_choice = CHOICE_RE.match(l2)
                if m_choice:
                    label = m_choice.group('label').strip()
                    target = m_choice.group('target').strip()
                    edges.setdefault(current_id, []).append((label, target))
                j += 1

            i = j
            continue
        i += 1

    return nodes, edges


def is_ending(node_id: str, title: str) -> bool:
    if any(node_id.startswith(p) for p in ENDING_ID_PREFIXES):
        return True
    return any(k in (title or '') for k in ENDING_TITLE_KEYWORDS)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--input', required=True)
    ap.add_argument('--output', required=True)
    args = ap.parse_args()

    nodes, edges = parse_story(args.input)

    # Compute out-degree
    outdeg = {nid: len(edges.get(nid, [])) for nid in nodes.keys()}

    leaf_nodes = [nid for nid, deg in outdeg.items() if deg == 0]

    endings = []
    non_endings = []
    for nid in sorted(leaf_nodes):
        title = nodes.get(nid, '')
        (endings if is_ending(nid, title) else non_endings).append((nid, title))

    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    with open(args.output, 'w', encoding='utf-8') as f:
        f.write('# 无后续（无 choices）节点清单\n\n')
        f.write(f'- 总节点数: {len(nodes)}\n')
        f.write(f'- 无后续节点数: {len(leaf_nodes)}\n')
        f.write(f'- 结局型: {len(endings)}\n')
        f.write(f'- 可能未接续: {len(non_endings)}\n\n')

        f.write('## 结局型（Leaf Endings）\n')
        for nid, title in endings:
            f.write(f'- {title} | id: `{nid}`\n')
        f.write('\n')

        f.write('## 可能未接续（需人工确认是否应有后续）\n')
        for nid, title in non_endings:
            f.write(f'- {title} | id: `{nid}`\n')

    print(f"Wrote report to {args.output}. Endings={len(endings)} NonEndings={len(non_endings)}")


if __name__ == '__main__':
    main()