#!/usr/bin/env python3
"""Align and alias the Pokedex voice stream inside the Radium sound export."""

from __future__ import annotations

import argparse
import csv
import difflib
import os
import re
import unicodedata
from pathlib import Path


FIELDS = [
    "pokemon_number",
    "pokemon_name",
    "role",
    "role_variant",
    "object_id",
    "section8_index",
    "source_wav_path",
    "alias_wav_path",
    "duration_seconds",
    "channels",
    "sample_rate",
    "command_ref_count",
    "alignment_action",
    "alignment_score",
    "transcript",
    "csv_category",
    "csv_text",
    "link_type",
]


def read_tsv(path: Path) -> list[dict[str, str]]:
    with path.open(newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f, delimiter="\t"))


def read_pokemon_rows(path: Path) -> list[dict[str, str]]:
    with path.open(newline="", encoding="utf-8", errors="replace") as f:
        rows = list(csv.DictReader(f))
    return sorted(
        [row for row in rows if row.get("in_game") == "TRUE" and row.get("name") != "Blank"],
        key=lambda row: int(row["number"]),
    )


def normalize(value: str) -> str:
    text = unicodedata.normalize("NFKD", value or "").encode("ascii", "ignore").decode("ascii")
    text = re.sub(r"[^a-zA-Z0-9]+", " ", text).lower()
    return re.sub(r"\s+", " ", text).strip()


def similarity(transcript: str, expected: str, role: str) -> float:
    got = normalize(transcript)
    want = normalize(expected)
    if not got or not want:
        return 0.0
    if role == "name":
        return difflib.SequenceMatcher(None, got.replace(" ", ""), want.replace(" ", "")).ratio()
    if role == "category":
        got = got.replace("the ", "").replace("pokemon", "").strip()
        want = want.replace("pokemon", "").strip()
    return difflib.SequenceMatcher(None, got[:180], want[:180]).ratio()


def slug(value: str, max_len: int = 72) -> str:
    text = unicodedata.normalize("NFKD", value).encode("ascii", "ignore").decode("ascii")
    text = text.upper().replace("&", " AND ")
    text = re.sub(r"[^A-Z0-9]+", "_", text).strip("_")
    text = re.sub(r"_+", "_", text)
    return (text[:max_len].rstrip("_") or "SOUND")


def link_alias(source_path: Path, alias_path: Path) -> str:
    alias_path.parent.mkdir(parents=True, exist_ok=True)
    if alias_path.exists() or alias_path.is_symlink():
        try:
            if alias_path.samefile(source_path):
                return "hardlink" if not alias_path.is_symlink() else "symlink"
        except OSError:
            pass
        return "existing"
    try:
        os.link(source_path, alias_path)
        return "hardlink"
    except OSError:
        alias_path.symlink_to(Path("..") / "sounds" / source_path.name)
        return "symlink"


def build_field_stream(pokemon_rows: list[dict[str, str]]) -> list[tuple[dict[str, str], str, str]]:
    fields: list[tuple[dict[str, str], str, str]] = []
    for row in pokemon_rows:
        fields.append((row, "description", row["description"]))
        fields.append((row, "name", row["name"]))
        fields.append((row, "category", row["category"]))
    return fields


def transition_cost(transcript: str, duration: float, text: str, role: str, duplicate: bool = False) -> tuple[float, float]:
    score = similarity(transcript, text, role)
    penalty = 0.0
    if role == "description" and duration < 2.0:
        penalty += 0.25
    if role == "name" and duration > 2.5:
        penalty += 0.45
    if role == "category" and duration > 3.5:
        penalty += 0.35
    if duplicate:
        penalty += 0.08
    return (1.0 - score) + penalty, score


def align_stream(
    object_ids: list[int],
    fields: list[tuple[dict[str, str], str, str]],
    asr_by_object: dict[int, dict[str, str]],
) -> list[tuple[int, int, str, float]]:
    states: dict[int, tuple[float, list[tuple[int, int, str, float]]]] = {0: (0.0, [])}
    for index, object_id in enumerate(object_ids):
        asr = asr_by_object[object_id]
        transcript = asr.get("transcript", "")
        duration = float(asr["duration_seconds"])
        next_states: dict[int, tuple[float, list[tuple[int, int, str, float]]]] = {}
        for field_index, (cost, path) in states.items():
            if field_index < len(fields):
                _, role, text = fields[field_index]
                delta, score = transition_cost(transcript, duration, text, role)
                new_index = field_index + 1
                candidate = (cost + delta, path + [(object_id, field_index, "match", score)])
                if new_index not in next_states or candidate[0] < next_states[new_index][0]:
                    next_states[new_index] = candidate
            if field_index > 0 and fields[field_index - 1][1] == "description":
                _, role, text = fields[field_index - 1]
                delta, score = transition_cost(transcript, duration, text, role, duplicate=True)
                candidate = (cost + delta, path + [(object_id, field_index - 1, "duplicate_description", score)])
                if field_index not in next_states or candidate[0] < next_states[field_index][0]:
                    next_states[field_index] = candidate

        min_index = max(0, index - 10)
        max_index = min(len(fields), index + 2)
        next_states = {key: value for key, value in next_states.items() if min_index <= key <= max_index}
        if len(next_states) > 50:
            next_states = dict(sorted(next_states.items(), key=lambda item: item[1][0])[:50])
        states = next_states

    final_index = len(fields)
    if final_index not in states:
        raise RuntimeError(f"alignment did not reach final field index {final_index}; candidates={sorted(states)}")
    return states[final_index][1]


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--pokemon-data", type=Path, default=Path("assets/godot_raw/data/pokemon_data.txt"))
    parser.add_argument("--sound-manifest", type=Path, default=Path("analysis/radium-image/sound_export_manifest.tsv"))
    parser.add_argument("--asr-manifest", type=Path, default=Path("analysis/radium-image/asr_draft_sound_labels.tsv"))
    parser.add_argument("--alias-dir", type=Path, default=Path("analysis/radium-image/sounds-pokedex-voice"))
    parser.add_argument("--out", type=Path, default=Path("analysis/radium-image/pokedex_voice_label_manifest.tsv"))
    parser.add_argument("--start-object", type=int, default=409)
    parser.add_argument("--end-object", type=int, default=1545)
    args = parser.parse_args()

    pokemon_rows = read_pokemon_rows(args.pokemon_data)
    fields = build_field_stream(pokemon_rows)
    sounds_by_object = {
        int(row["primary_object_local_id"]): row
        for row in read_tsv(args.sound_manifest)
        if row.get("primary_object_local_id", "").isdigit()
    }
    asr_by_object = {int(row["object_id"]): row for row in read_tsv(args.asr_manifest) if row["object_id"].isdigit()}
    object_ids = list(range(args.start_object, args.end_object + 1))
    missing = [object_id for object_id in object_ids if object_id not in sounds_by_object or object_id not in asr_by_object]
    if missing:
        raise RuntimeError(f"missing sound/ASR rows for objects: {missing[:20]}")

    path = align_stream(object_ids, fields, asr_by_object)
    variant_counts: dict[tuple[str, str], int] = {}
    rows: list[dict[str, object]] = []
    for object_id, field_index, action, score in path:
        pokemon, role, text = fields[field_index]
        key = (pokemon["number"], role)
        variant_counts[key] = variant_counts.get(key, 0) + 1
        role_variant = variant_counts[key]
        sound = sounds_by_object[object_id]
        asr = asr_by_object[object_id]
        role_label = role.upper()
        if role == "category":
            role_label = f"CATEGORY_{slug(pokemon['category'], 36)}"
        elif action == "duplicate_description" or role_variant > 1:
            role_label = f"DESCRIPTION_ALT{role_variant}"
        alias_name = (
            f"POKEDEX_{int(pokemon['number']):04d}_{slug(pokemon['name'], 36)}_{role_label}"
            f"__obj{object_id:04d}_s{int(sound['section8_index']):04d}"
            f"_{sound['channels']}ch_{sound['sample_rate']}hz.wav"
        )
        alias_path = args.alias_dir / alias_name
        source_path = Path(sound["wav_path"])
        rows.append(
            {
                "pokemon_number": pokemon["number"],
                "pokemon_name": pokemon["name"],
                "role": role,
                "role_variant": role_variant,
                "object_id": object_id,
                "section8_index": sound["section8_index"],
                "source_wav_path": sound["wav_path"],
                "alias_wav_path": alias_path.as_posix(),
                "duration_seconds": sound["duration_seconds"],
                "channels": sound["channels"],
                "sample_rate": sound["sample_rate"],
                "command_ref_count": sound["command_ref_count"],
                "alignment_action": action,
                "alignment_score": f"{score:.6f}",
                "transcript": asr.get("transcript", ""),
                "csv_category": pokemon["category"],
                "csv_text": text,
                "link_type": link_alias(source_path, alias_path),
            }
        )

    args.out.parent.mkdir(parents=True, exist_ok=True)
    with args.out.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=FIELDS, delimiter="\t")
        writer.writeheader()
        writer.writerows(rows)
    duplicate_rows = [row for row in rows if row["alignment_action"] == "duplicate_description"]
    print(
        f"wrote {args.out} rows={len(rows)} pokemon={len(pokemon_rows)} "
        f"duplicates={len(duplicate_rows)} alias_dir={args.alias_dir}"
    )
    for row in duplicate_rows:
        print(f"duplicate_description obj={row['object_id']} pokemon={row['pokemon_number']} {row['pokemon_name']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
