"""plan.py — timing.json → plan.json scene planner.

CLI:
    py -3.14 scripts/plan.py <script.txt> <timing.json> <out_plan.json>

Input timing.json schema:
    {
        "voice": "...",
        "audio_duration": float,
        "words": [{"text": str, "start": float, "dur": float}, ...]
    }

Output plan.json schema: see SPEC in module docstring.
"""

from __future__ import annotations

import json
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------

@dataclass
class Word:
    """A single word boundary entry from timing.json."""

    text: str
    start: float
    dur: float

    @property
    def end(self) -> float:
        return self.start + self.dur


@dataclass
class KeywordHit:
    """A keyword found inside a scene with its absolute timestamp."""

    text: str
    t: float


@dataclass
class Scene:
    """One scene in the output plan."""

    id: int
    text: str
    start: float
    end: float
    intent: str
    pose: str
    face: str
    keywords: list[KeywordHit]
    subtitle: str
    prop: Optional[str]

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "text": self.text,
            "start": self.start,
            "end": self.end,
            "intent": self.intent,
            "pose": self.pose,
            "face": self.face,
            "keywords": [{"text": kw.text, "t": kw.t} for kw in self.keywords],
            "subtitle": self.subtitle,
            "prop": self.prop,
        }


# ---------------------------------------------------------------------------
# Step 1 – sentence segmentation
# ---------------------------------------------------------------------------

# Sentence-ending punctuation (full-width and ASCII variants).
_SENT_END = re.compile(r"[。！？!?；;]+")


def split_sentences(script_text: str) -> list[str]:
    """Split *script_text* into sentences on 。！？!?；; boundaries.

    Empty strings and pure-whitespace strings are discarded.

    Args:
        script_text: Raw multi-line script content.

    Returns:
        Ordered list of sentence strings, each retaining its trailing
        punctuation.
    """
    # Collapse newlines / extra spaces first.
    text = script_text.replace("\r\n", "\n").replace("\r", "\n")
    # Tokenise: split AFTER each sentence-ending punctuation cluster.
    parts = _SENT_END.split(text)
    # Because split() removes the delimiter, we re-attach delimiters.
    delimiters = _SENT_END.findall(text)

    sentences: list[str] = []
    for i, part in enumerate(parts):
        chunk = part.strip()
        if chunk:
            if i < len(delimiters):
                chunk += delimiters[i]
            sentences.append(chunk)
        elif i < len(delimiters) and sentences:
            # Trailing punctuation with no body — attach to previous sentence.
            sentences[-1] += delimiters[i]

    return sentences


# ---------------------------------------------------------------------------
# Step 2 – word → sentence alignment (greedy, fault-tolerant)
# ---------------------------------------------------------------------------

_PUNCT_RE = re.compile(
    r"[\s　＀-￯‘-‟—…"
    r"，。！？；：、“”‘’"
    r"（）【】,.!?;:’\"\(\)\[\]{}\-]+"
)


def _strip_for_match(s: str) -> str:
    """Strip punctuation and whitespace for token matching."""
    return _PUNCT_RE.sub("", s)


def align_words_to_sentences(
    sentences: list[str], words: list[Word]
) -> list[list[Word]]:
    """Greedily assign each *word* to its sentence.

    Strategy:
    - Build a clean (no-punct, no-space) version of each sentence.
    - Maintain a pointer into the current sentence's clean text.
    - For each word token, try to find it at (or near) the pointer.
    - On match: advance pointer by len(token).
    - On miss: fall through — still assign to current sentence (fault-tolerant).
    - When current sentence is exhausted advance to next sentence.

    Args:
        sentences: Ordered sentence strings.
        words: Ordered Word objects from timing.json.

    Returns:
        List of length ``len(sentences)``, each element being the list of
        Words assigned to that sentence.
    """
    # Pre-compute clean versions of sentences.
    clean_sents = [_strip_for_match(s) for s in sentences]

    buckets: list[list[Word]] = [[] for _ in sentences]
    sent_idx = 0
    sent_ptr = 0  # character pointer into clean_sents[sent_idx]

    for word in words:
        # Advance to next sentence if current one is exhausted.
        while sent_idx < len(sentences) - 1 and sent_ptr >= len(clean_sents[sent_idx]):
            sent_idx += 1
            sent_ptr = 0

        tok = _strip_for_match(word.text)
        if not tok:
            # Whitespace-only token — attach to current sentence, skip.
            buckets[sent_idx].append(word)
            continue

        # Look for the token starting at sent_ptr (allow a small lookahead
        # to skip over punctuation chars that may not appear in words).
        clean = clean_sents[sent_idx]
        found_at = clean.find(tok, sent_ptr)

        if found_at != -1:
            # Happy path: token found in current sentence.
            sent_ptr = found_at + len(tok)
            buckets[sent_idx].append(word)
        else:
            # Try remaining sentences (in order) to handle edge cases.
            matched_sent = -1
            matched_ptr = 0
            for future_idx in range(sent_idx + 1, len(sentences)):
                future_clean = clean_sents[future_idx]
                pos = future_clean.find(tok)
                if pos != -1:
                    matched_sent = future_idx
                    matched_ptr = pos + len(tok)
                    break

            if matched_sent != -1:
                # Token found in a future sentence — flush forward.
                sent_idx = matched_sent
                sent_ptr = matched_ptr
                buckets[sent_idx].append(word)
            else:
                # Totally unmatched: fault-tolerant fallback — keep in current.
                buckets[sent_idx].append(word)
                # Do NOT advance sent_ptr so next token still tries here.

    return buckets


# ---------------------------------------------------------------------------
# Step 3 – intent classification
# ---------------------------------------------------------------------------

# Priority order (highest first): impact > cheer > excite > discover > sad > walk > talk
_INTENT_RULES: list[tuple[str, re.Pattern]] = [
    ("impact",   re.compile(r"撞|摔|砰")),
    ("cheer",    re.compile(r"开心|高兴|蹦蹦跳跳|欢呼")),
    ("excite",   re.compile(r"！$|只要|竟然|居然")),
    ("discover", re.compile(r"发现|突然|原来")),
    ("sad",      re.compile(r"难过|哭|伤心")),
    ("walk",     re.compile(r"走|路|上学路上")),
    ("talk",     re.compile(r".")),  # default — always matches
]


def classify_intent(sentence: str) -> str:
    """Return the highest-priority intent that matches *sentence*.

    Args:
        sentence: A single sentence string (with trailing punctuation).

    Returns:
        Intent string: one of walk|discover|excite|cheer|sad|impact|talk.
    """
    for intent, pattern in _INTENT_RULES:
        if pattern.search(sentence):
            return intent
    return "talk"


# ---------------------------------------------------------------------------
# Step 4 – pose / face mapping
# ---------------------------------------------------------------------------

_POSE_FACE: dict[str, tuple[str, str]] = {
    "walk":     ("walk",  "calm"),
    "discover": ("idle",  "surprised"),
    "excite":   ("cheer", "surprised"),
    "cheer":    ("cheer", "calm"),
    "sad":      ("cower", "sad"),
    "impact":   ("cower", "surprised"),
    "talk":     ("idle",  "calm"),
}


def map_pose_face(intent: str) -> tuple[str, str]:
    """Return (pose, face) for a given *intent*.

    Args:
        intent: Intent string.

    Returns:
        Tuple of (pose, face) strings.
    """
    return _POSE_FACE.get(intent, ("idle", "calm"))


# ---------------------------------------------------------------------------
# Step 5 – keyword extraction (numeric quantifiers)
# ---------------------------------------------------------------------------

# Matches numeric quantifiers such as 三公里, 十分钟, 100米.
# Units are split into two groups:
#   - Unambiguous: 公里 米 分钟 小时 步 秒 — always valid measurement units.
#   - Ambiguous: 天 年 月 — also appear in idiomatic/temporal expressions.
# For ambiguous units we require a multi-char numeric prefix (十天, 三十天)
# OR an Arabic-digit prefix, which rules out bare 一天/一月/一年 idioms.
_KW_RE = re.compile(
    # Unambiguous units — any numeric prefix (including single digit).
    r"[零一二三四五六七八九十百千万亿两]+(?:公里|分钟|小时|步|秒|米(?!道))"
    r"|[0-9]+(?:公里|分钟|小时|步|秒|米(?!道))"
    # Ambiguous units — calendar words appear in idioms (一天/一月), and
    # money/count words appear in idioms too (一块去, 每一次), so require
    # a multi-char numeric prefix OR an Arabic-digit prefix.
    r"|[零一二三四五六七八九十百千万亿两]{2,}(?:[天年月]|块钱?|元|次|倍)"
    r"|[0-9]+(?:[天年月]|块钱?|元|次|倍)"
)


def extract_keywords(sentence: str, sent_words: list[Word]) -> list[KeywordHit]:
    """Find numeric-quantifier keywords in *sentence* and map them to word timestamps.

    For each keyword match, we find the Word whose text contains (or best
    overlaps) the keyword's characters, and use that word's start as ``t``.

    Args:
        sentence: Sentence text.
        sent_words: Words assigned to this sentence.

    Returns:
        List of KeywordHit objects.
    """
    # Concatenate the stripped word texts in order, remembering which word
    # each character position belongs to, so a keyword spanning several
    # tokens (e.g. "十分钟" split as "十"+"分钟") maps to the FIRST token's
    # start time — the moment the keyword begins to be spoken.
    concat_chars: list[str] = []
    char_owner: list[Word] = []
    for w in sent_words:
        for ch in _strip_for_match(w.text):
            concat_chars.append(ch)
            char_owner.append(w)
    concat = "".join(concat_chars)

    hits: list[KeywordHit] = []
    search_from = 0
    for m in _KW_RE.finditer(sentence):
        kw_text = m.group()
        kw_clean = _strip_for_match(kw_text)

        idx = concat.find(kw_clean, search_from)
        if idx < 0:
            idx = concat.find(kw_clean)

        best_word: Optional[Word] = None
        if idx >= 0:
            best_word = char_owner[idx]
            search_from = idx + len(kw_clean)
        else:
            # Fallback: the word with the most character overlap.
            kw_chars = set(kw_clean)
            best_overlap = 0
            for w in sent_words:
                overlap = len(kw_chars & set(_strip_for_match(w.text)))
                if overlap > best_overlap:
                    best_overlap = overlap
                    best_word = w

        if best_word is not None:
            hits.append(KeywordHit(text=kw_text, t=best_word.start))

    return hits


# ---------------------------------------------------------------------------
# Step 6 – prop classification
# ---------------------------------------------------------------------------

_PROP_RULES: list[tuple[re.Pattern, str]] = [
    (re.compile(r"路|山路|小路"), "path"),
    (re.compile(r"学校"),         "school"),
]


def classify_prop(sentence: str) -> Optional[str]:
    """Return a prop string or None.

    Rules are evaluated in order; first match wins.

    Args:
        sentence: Sentence text.

    Returns:
        Prop string or None.
    """
    for pattern, prop in _PROP_RULES:
        if pattern.search(sentence):
            return prop
    return None


# ---------------------------------------------------------------------------
# Main pipeline
# ---------------------------------------------------------------------------

def build_plan(script_text: str, timing: dict) -> dict:
    """Build a plan dict from *script_text* and *timing* data.

    Args:
        script_text: Raw script content.
        timing: Parsed timing.json dict.

    Returns:
        Plan dict ready for JSON serialisation.
    """
    words = [Word(w["text"], w["start"], w["dur"]) for w in timing["words"]]
    sentences = split_sentences(script_text)

    if not sentences:
        return {"scenes": []}

    buckets = align_words_to_sentences(sentences, words)

    scenes: list[Scene] = []
    for i, (sent, sent_words) in enumerate(zip(sentences, buckets)):
        if sent_words:
            start = sent_words[0].start
            end = sent_words[-1].start + sent_words[-1].dur
        elif scenes:
            # No words aligned — use previous scene's end as a stub.
            start = scenes[-1].end
            end = start + 0.5
        else:
            start = 0.0
            end = 0.5

        intent = classify_intent(sent)
        pose, face = map_pose_face(intent)
        keywords = extract_keywords(sent, sent_words)
        prop = classify_prop(sent)

        scenes.append(Scene(
            id=i + 1,
            text=sent,
            start=start,
            end=end,
            intent=intent,
            pose=pose,
            face=face,
            keywords=keywords,
            subtitle=sent,
            prop=prop,
        ))

    return {"scenes": [s.to_dict() for s in scenes]}


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

def main(argv: list[str] | None = None) -> int:
    """CLI entry point.

    Args:
        argv: Argument list (defaults to sys.argv[1:]).

    Returns:
        Exit code (0 = success, non-zero = error).
    """
    args = argv if argv is not None else sys.argv[1:]
    if len(args) != 3:
        print(
            "Usage: py -3.14 scripts/plan.py <script.txt> <timing.json> <out_plan.json>",
            file=sys.stderr,
        )
        return 1

    script_path, timing_path, out_path = Path(args[0]), Path(args[1]), Path(args[2])

    if not script_path.exists():
        print(f"ERROR: script file not found: {script_path}", file=sys.stderr)
        return 1
    if not timing_path.exists():
        print(f"ERROR: timing file not found: {timing_path}", file=sys.stderr)
        return 1

    script_text = script_path.read_text(encoding="utf-8")
    timing = json.loads(timing_path.read_text(encoding="utf-8"))

    plan = build_plan(script_text, timing)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(plan, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Wrote {len(plan['scenes'])} scenes → {out_path}")
    return 0


if __name__ == "__main__":
    sys.exit(main())
