import json
import os
import re
from urllib.parse import quote

def is_kanji(ch):
    return '\u4e00' <= ch <= '\u9faf'

def make_ruby(text, reading):
    if not text or not reading or text == reading:
        return text
    
    # Exclude ＊ from ruby
    prefix = ""
    if text.startswith('＊'):
        prefix = '＊'
        text = text[1:]
    
    # Matching tail
    i, j = len(text) - 1, len(reading) - 1
    while i >= 0 and j >= 0 and text[i] == reading[j]:
        i, j = i - 1, j - 1
    
    tail = text[i+1:]
    rem_text, rem_reading = text[:i+1], reading[:j+1]
    
    # Matching head
    i, j = 0, 0
    while i < len(rem_text) and j < len(rem_reading) and rem_text[i] == rem_reading[j]:
        i, j = i + 1, j + 1
    
    head = rem_text[:i]
    rem_text, rem_reading = rem_text[i:], rem_reading[j:]
    
    if not rem_text:
        return head + tail
    
    if all(is_kanji(c) for c in rem_text):
        return f"{prefix}{head}<ruby>{rem_text}<rt>{rem_reading}</rt></ruby>{tail}"
    
    match = re.search(r'[ぁ-んァ-ヶ]', rem_text)
    if match:
        kana_in_text = match.group()
        pos = rem_reading.find(kana_in_text)
        if pos != -1:
            k1, r1 = rem_text[:match.start()], rem_reading[:pos]
            k2, r2 = rem_text[match.end():], rem_reading[pos+1:]
            res = head
            if k1: res += f"<ruby>{k1}<rt>{r1}</rt></ruby>"
            res += kana_in_text
            if k2: res += make_ruby(k2, r2)
            return prefix + res + tail

    return prefix + f"{head}<ruby>{rem_text}<rt>{rem_reading}</rt></ruby>{tail}"

def fix_internal_links(html, label_to_hw_map):
    if not html: return html
    
    # Pass 1: Handle links with <sup> (target labels like 400, 105a etc)
    def replace_sup_link(match):
        full_tag = match.group(0)
        # Extract the label inside <sup>
        sup_match = re.search(r'<sup[^>]*>\s*([^<]+)\s*</sup>', full_tag, re.IGNORECASE)
        if sup_match:
            label = sup_match.group(1).strip()
            target_hw = label_to_hw_map.get(label)
            if target_hw:
                # Replace the href inside the current match
                return re.sub(r'href="[^"]+"', f'href="entry://{target_hw}"', full_tag)
        return full_tag

    # Pass 1: Resolve internal links with <sup>
    html = re.sub(r'<a [^>]*>.*?<sup[^>]*>.*?</sup>.*?</a>', replace_sup_link, html, flags=re.IGNORECASE | re.DOTALL)

    # Pass 2: Handle links without <sup> but with labels in text (e.g., <a href="02066">⇒104c</a>)
    def replace_text_label_link(match):
        full_tag = match.group(0)
        # Extract text content (strip tags)
        text_content = re.sub(r'<[^>]+>', '', full_tag)
        # Look for labels like 061b, 104c, 400 in text content
        # Patterns: 1-3 digits + optional letter (a-z)
        label_match = re.search(r'([a-z]?\d{1,4}[a-z]?)', text_content, re.IGNORECASE)
        if label_match:
            label = label_match.group(1).strip()
            target_hw = label_to_hw_map.get(label)
            if target_hw:
                return re.sub(r'href="[^"]+"', f'href="entry://{target_hw}"', full_tag)
        return full_tag

    # Pass 2: Resolve internal links based on text content labels
    html = re.sub(r'<a [^>]*>.*?</a>', replace_text_label_link, html, flags=re.IGNORECASE | re.DOTALL)
    
    # Step 2: Extract block sections and wrap them
    # We find all occurrences of markers and split the string
    markers = [r'○', r'〔注〕', r'〔参考〕']
    # Create a regex to split by any of these markers, keeping the marker
    pattern = '|'.join(map(re.escape, markers))
    
    parts = re.split(f'({pattern})', html)
    
    new_html = parts[0] # Text before first marker
    for i in range(1, len(parts), 2):
        marker = parts[i]
        content = parts[i+1] if i+1 < len(parts) else ""
        content.strip('　')
        
        # Determine class based on marker
        cls = "bullet-tag"
        if marker == "〔注〕": cls = "note-tag"
        elif marker == "〔参考〕": cls = "ref-tag"
        
        new_html += f'<div class="block-section"><span class="{cls}">{marker}</span>{content}</div>'
    
    # Step 3: Wrap inline usage tags in the resulting HTML
    # (Exclude markers we already handled)
    new_html = re.sub(r'[\u3000\s]*(〔(?!(注|参考))[^〕]+〕)', r'<span class="usage-tag">\1</span>', new_html)
    
    return new_html

def get_headword(node):
    if not node: return None
    node_idx = node.get('_nodeIndex')
    if node_idx == 0: return "kadokawaruigo_index"
    label = node.get('_label', '')
    title = node.get('_title', '')
    if not label and not title: return None
    if not label: return title
    if not title: return label
    return f"{label} {title}"

def process_thesaurus():
    print("Loading data...")
    with open('index.json', 'r', encoding='utf-8') as f:
        index_data = json.load(f)
    with open('contents.json', 'r', encoding='utf-8') as f:
        contents_data = json.load(f)
    
    # Prep maps
    print("Preparing maps...")
    label_to_hw = {}
    for node in index_data.values():
        if node.get('_kind') in [1, 2]:
            label = node.get('_label')
            if label:
                label_to_hw[label] = get_headword(node)

    # Step 1: Pre-process Kind 4 "sub-category" capturing
    print("Grouping sub-categories (Kind 4)...")
    for main_node in index_data.values():
        child_indices = main_node.get('_childIndexes', [])
        if not child_indices: continue
        
        current_kind4 = None
        main_node['_captured_by_kind4'] = set()
        
        for idx in child_indices:
            child = index_data.get(str(idx))
            if not child: continue
            
            kind = child.get('_kind')
            if kind == 4:
                current_kind4 = child
                if '_own_captured_words' not in current_kind4:
                    current_kind4['_own_captured_words'] = []
            elif kind == 3:
                # Capture Word into previous Kind 4 node if it exists
                if current_kind4:
                    current_kind4['_own_captured_words'].append(idx)
                    main_node['_captured_by_kind4'].add(idx)
            else:
                # Kind 1, 2, etc. (Real sub-categories) - stop capturing
                current_kind4 = None

    # Step 2: Merge categories with same Headword
    hw_to_nodes = {}
    for node_id, node in index_data.items():
        # Add content
        c_id = str(node.get('_contentsID', -1))
        node['content'] = contents_data.get(c_id) if c_id != "-1" else None
        
        if node.get('_kind') in [0, 1, 2, 4] or node.get('_childIndexes'):
            hw = get_headword(node)
            if hw:
                if hw not in hw_to_nodes: hw_to_nodes[hw] = []
                # Avoid duplicate nodes in the same headword entry
                if node not in hw_to_nodes[hw]:
                    hw_to_nodes[hw].append(node)

    print(f"Generating MDX...")
    mdx_file = 'KADOKAWARUIGO_final.mdx.txt'
    
    def get_breadcrumbs(node, all_nodes):
        path = [node]
        curr = node
        seen_idxs = set()
        while curr:
            n_idx = curr.get('_nodeIndex')
            if n_idx in seen_idxs: break
            seen_idxs.add(n_idx)
            
            p_idx = curr.get('_parentNodeIndex')
            if p_idx == -1:
                if curr.get('_nodeIndex') == 0: path.append(curr)
                break
            parent = all_nodes.get(str(p_idx))
            if parent:
                if parent.get('_kind') in [0, 1, 2, 4]: path.append(parent)
                curr = parent
            else: break
        
        # Filter to avoid "kadokawaruigo_index" being added twice 
        # (if it's already in path through parent chain)
        res = []
        seen_hws = set()
        for p in reversed(path):
            hw = get_headword(p)
            if hw not in seen_hws:
                res.append(p)
                seen_hws.add(hw)
        return res

    with open(mdx_file, 'w', encoding='utf-8') as f:
        # Special Entry #index redirection
        f.write("#index\n@@@LINK=kadokawaruigo_index\n</>\n")

        # Special Entry #hanrei
        if os.path.exists('usage/hanrei.html'):
            with open('usage/hanrei.html', 'r', encoding='utf-8') as hf:
                h_body = re.search(r'<body>(.*?)</body>', hf.read(), re.DOTALL | re.IGNORECASE)
                if h_body:
                    f.write("#hanrei\n")
                    f.write('<link rel="stylesheet" type="text/css" href="KADOKAWARUIGO.css">\n')
                    f.write('<div class="thesaurus-container hanrei-entry">\n')
                    f.write(h_body.group(1).strip() + '\n')
                    f.write('</div>\n</>\n')

        # Redirection links for Words (Leafs)
        # We handle them separately first to avoid headword collisions with merged categories
        for node in index_data.values():
            if not node.get('_childIndexes') and node.get('_kind') == 3:
                hw_title = node.get('_title')
                hw_yomi = node.get('_yomi')
                
                # Strip ＊ from word title/yomi for redirection entry
                if hw_title: hw_title = hw_title.lstrip('＊')
                if hw_yomi: hw_yomi = hw_yomi.lstrip('＊')
                
                p_idx = node.get('_parentNodeIndex')
                parent = index_data.get(str(p_idx))
                if parent:
                    p_hw = get_headword(parent)
                    if p_hw:
                        if hw_title:
                            f.write(f"{hw_title}\n@@@LINK={p_hw}\n</>\n")
                        if hw_yomi and hw_yomi != hw_title:
                            f.write(f"{hw_yomi}\n@@@LINK={p_hw}\n</>\n")

        # Merged Categories
        # Sort by first node's index
        sorted_hws = sorted(hw_to_nodes.keys(), key=lambda x: hw_to_nodes[x][0]['_nodeIndex'])
        
        for hw in sorted_hws:
            nodes = hw_to_nodes[hw]
            primary = nodes[0]
            node_idx = primary.get('_nodeIndex')
            
            f.write(f"{hw}\n")
            f.write('<link rel="stylesheet" type="text/css" href="KADOKAWARUIGO.css">\n')
            f.write('<div class="thesaurus-container">\n')
            
            # Nav
            if node_idx != 0:
                f.write('  <div class="breadcrumb-nav">\n')
                crumbs = get_breadcrumbs(primary, index_data)
                # crumbs[:-1] because current entry is the last one in the breadcrumb list
                links = [f'<a href="entry://{get_headword(c)}">{c.get("_title") or "TOP"}</a>' for c in crumbs[:-1]]
                f.write(f'    {" <span class=\"sep\">›</span> ".join(links)}\n')
                f.write('  </div>\n')
            
            # Combined Header
            f.write('  <div class="category-header">\n')
            if node_idx == 0:
                f.write('    <span class="title">角川类语新词典</span>\n')
            else:
                f.write(f'    <span class="label">{primary.get("_label", "")}</span>\n')
                title, yomi = primary.get("_title", ""), primary.get("_yomi", "")
                ruby_title = make_ruby(title, yomi) if yomi and yomi != title else title
                f.write(f'    <span class="title">{ruby_title}</span>\n')
            f.write('  </div>\n')
            
            # Combine all captions and explanations from merged nodes
            seen_captions = set()
            for n in nodes:
                cap = n.get('_caption')
                if cap and cap not in seen_captions:
                    f.write(f'  <div class="caption">{cap}</div>\n')
                    seen_captions.add(cap)
            
            seen_explanations = set()
            for n in nodes:
                exp = n.get('content')
                if exp and exp not in seen_explanations:
                    exp = fix_internal_links(exp, label_to_hw)
                    f.write(f'  <div class="explanation">{exp}</div>\n')
                    seen_explanations.add(exp)
            
            # Combine all children
            all_cat_children = []
            all_word_children = []
            seen_child_indices = set()
            for n in nodes:
                # 1. Normal children
                child_indices = n.get('_childIndexes', [])
                # 2. Captured words (specifically for Kind 4 nodes)
                captured_indices = n.get('_own_captured_words', [])
                
                # We skip words that have been captured by sub-categories in this parent
                excluded_indices = n.get('_captured_by_kind4', set())
                
                for idx in child_indices + captured_indices:
                    if idx in seen_child_indices: continue
                    if idx in excluded_indices: continue
                    
                    seen_child_indices.add(idx)
                    child = index_data.get(str(idx))
                    if not child: continue
                    
                    if child.get('_childIndexes') or child.get('_kind') in [1, 2, 4]:
                        all_cat_children.append(child)
                    else:
                        all_word_children.append(child)
            
            if all_cat_children or all_word_children:
                f.write('  <div class="children-list">\n')
                for child in all_cat_children:
                    c_hw = get_headword(child)
                    if c_hw == hw: continue # Skip self-links
                    title, cap = child.get('_title', ''), child.get('_caption', '')
                    item = f'<a href="entry://{c_hw}">{get_headword(child)}</a>'
                    if cap: item += f' <span class="child-caption">{cap}</span>'
                    f.write(f'    <div class="child-item">{item}</div>\n')
                
                if all_word_children:
                    for child in all_word_children:
                        f.write('    <div class="word-entry">\n')
                        if child.get('_label'): f.write(f' <span class="label">{child.get("_label")}</span>\n')
                        wt, wy = child.get("_title", ""), child.get("_yomi", "")
                        # Make sure ＊ is passed to make_ruby
                        f.write(f' <span class="word-title">{make_ruby(wt, wy) if wy and wy != wt else wt}</span>\n')
                        exp = child.get('content')
                        if exp: f.write(f' <div class="word-explanation">{fix_internal_links(exp, label_to_hw)}</div>\n')
                        f.write('    </div>\n')
                f.write('  </div>\n')
            f.write('</div>\n</>\n')

    print(f"MDX saved to {mdx_file}")

if __name__ == "__main__":
    process_thesaurus()