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

from __future__ import annotations

import argparse
import html
import json
import os
import re
import sys
from dataclasses import dataclass, asdict
from typing import Dict, List, Any

from fugashi import Tagger


# Ensure UTF-8 output
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="backslashreplace")


KATAKANA_START = ord("ァ")
KATAKANA_END = ord("ヶ")
KATAKANA_TO_HIRAGANA_OFFSET = ord("ぁ") - ord("ァ")


def katakana_to_hiragana(text: str) -> str:
    return "".join(
        chr(ord(ch) + KATAKANA_TO_HIRAGANA_OFFSET) if KATAKANA_START <= ord(ch) <= KATAKANA_END else ch
        for ch in text
    )


def has_japanese(text: str) -> bool:
    return any(
        0x3040 <= ord(ch) <= 0x30FF or
        0x3400 <= ord(ch) <= 0x9FFF
        for ch in text
    )


def is_content_word(pos1: str) -> bool:
    return pos1 in {"名詞", "動詞", "形容詞", "副詞"}


def should_skip(surface: str, pos1: str) -> bool:
    return not surface.strip() or pos1 == "補助記号"


def clean_lemma(lemma: str) -> str:
    return re.split(r"-", lemma, maxsplit=1)[0] if lemma else lemma


def get_feature(feature, name: str, default=""):
    return getattr(feature, name, default) or default


def build_ruby(surface: str, reading: str) -> str:
    if not reading or surface == reading or not has_japanese(surface):
        return html.escape(surface)
    return f"<ruby>{html.escape(surface)}<rt>{html.escape(reading)}</rt></ruby>"


@dataclass
class TokenReading:
    surface: str
    reading_hiragana: str
    lemma: str
    lemma_reading_hiragana: str
    pos: List[str]


# ---------- Persistent vocab handling ----------

def load_vocab(path: str) -> Dict[str, Dict]:
    if not os.path.exists(path):
        return {}

    with open(path, "r", encoding="utf-8") as f:
        raw = json.load(f)

    # normalize + convert surfaces_seen → set
    out = {}
    for lemma, entry in raw.items():
        out[lemma] = {
            "lemma": entry.get("lemma", lemma),
            "lemma_reading_hiragana": entry.get("lemma_reading_hiragana", ""),
            "pos": entry.get("pos", []),
            "mention_count": int(entry.get("mention_count", 0)),
            "surfaces_seen": set(entry.get("surfaces_seen", [])),  # ← key change
        }
    return out


def save_vocab(path: str, vocab: Dict[str, Dict]):
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)

    # convert sets → sorted lists
    serializable = {
        lemma: {
            **entry,
            "surfaces_seen": sorted(entry.get("surfaces_seen", []))
        }
        for lemma, entry in vocab.items()
    }

    with open(path, "w", encoding="utf-8") as f:
        json.dump(serializable, f, ensure_ascii=False, indent=2)


def update_vocab(vocab: Dict[str, Dict], tokens: List[TokenReading]):
    for t in tokens:
        pos1 = t.pos[0] if t.pos else ""

        if should_skip(t.surface, pos1) or not is_content_word(pos1):
            continue

        entry = vocab.get(t.lemma)

        if entry is None:
            vocab[t.lemma] = {
                "lemma": t.lemma,
                "lemma_reading_hiragana": t.lemma_reading_hiragana,
                "pos": t.pos,
                "mention_count": 1,
                "surfaces_seen": {t.surface},  # ← set
            }
        else:
            entry["mention_count"] += 1
            entry.setdefault("surfaces_seen", set()).add(t.surface)


# ---------- Main analysis ----------

def analyze(text: str, vocab: Dict[str, Dict]):
    tagger = Tagger()

    tokens: List[TokenReading] = []
    ruby_parts = []

    for word in tagger(text):
        f = word.feature
        surface = word.surface

        pos = [
            get_feature(f, "pos1"),
            get_feature(f, "pos2"),
            get_feature(f, "pos3"),
            get_feature(f, "pos4"),
        ]

        kana = get_feature(f, "kana")
        reading = katakana_to_hiragana(kana) if kana else ""

        lemma = clean_lemma(get_feature(f, "lemma") or surface)

        lform = get_feature(f, "lForm")
        kana_base = get_feature(f, "kanaBase")
        lemma_reading = katakana_to_hiragana(lform or kana_base or kana or lemma)

        token = TokenReading(surface, reading, lemma, lemma_reading, pos)
        tokens.append(token)

        ruby_parts.append(build_ruby(surface, reading))

    update_vocab(vocab, tokens)

    return {
        "original_text": text,
        "ruby_html": "".join(ruby_parts),
        "token_readings": [asdict(t) for t in tokens],
    }


# ---------- CLI ----------

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--vocab-file", required=True)
    parser.add_argument("text", nargs="*")
    args = parser.parse_args()

    if args.text:
        text = " ".join(args.text)
    else:
        text = sys.stdin.read()

    vocab = load_vocab(args.vocab_file)
    result = analyze(text, vocab)
    save_vocab(args.vocab_file, vocab)

    result["vocab_file"] = args.vocab_file
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()