"""Known-transcript alignment, separate from synthesis and its GPU lease."""
import logging
import math
from pathlib import Path
import unicodedata

import spaces  # Must precede torch.
import numpy as np
import torch
import torchaudio.functional as AF
from transformers import AutoModelForTokenClassification, AutoProcessor

from config import ALIGNER_ID, ALIGNER_REVISION, SAMPLE_RATE
from lifecycle import check_cancelled

log = logging.getLogger("zonos.alignment")
ALIGNER = None


def canonical(text):
    # Keep a source-index map so punctuation and typography survive unchanged.
    chars, positions = [], []
    for i, char in enumerate(text):
        for normalized in unicodedata.normalize("NFKC", char).casefold():
            if normalized.isalnum():
                chars.append(normalized)
                positions.append(i)
    return "".join(chars), positions


def attach_timings(text, words, duration):
    """Validate measured timestamps and map them back to the exact input text."""
    stream, positions = canonical(text)
    if not stream or not words or not math.isfinite(duration) or duration <= 0:
        raise ValueError("Missing alignment")
    cursor, previous, segments = 0, 0.0, []
    for word in words:
        token, _ = canonical(str(word["text"]))
        start, end = float(word["start_time"]), float(word["end_time"])
        if not token or not stream.startswith(token, cursor):
            raise ValueError("Alignment transcript mismatch")
        if not all(math.isfinite(v) for v in (start,end)) or start < previous-.001 or end < start or start < -.08 or end > duration+.16:
            raise ValueError("Invalid alignment timestamps")
        start, end = max(0,start), min(duration,end)
        if start > end:
            raise ValueError("Alignment outside audio")
        offset = positions[cursor]
        segments.append({"offset":offset,"start":round(start,3),"end":round(end,3)})
        cursor += len(token)
        previous = end
    if cursor != len(stream) or not any(w["end"] > w["start"] for w in segments):
        raise ValueError("Incomplete alignment")
    # Text strings, not offsets, cross the Python/JS boundary (Unicode-safe).
    result=[]
    for i, segment in enumerate(segments):
        left = 0 if i == 0 else segment["offset"]
        right = segments[i+1]["offset"] if i+1<len(segments) else len(text)
        result.append({"text":text[left:right],"start":segment["start"],"end":segment["end"]})
    return result


class Aligner:
    def __init__(self):
        log.info("alignment_loading revision=%s", ALIGNER_REVISION)
        self.processor=AutoProcessor.from_pretrained(ALIGNER_ID,revision=ALIGNER_REVISION)
        self.model=AutoModelForTokenClassification.from_pretrained(
            ALIGNER_ID, revision=ALIGNER_REVISION, dtype=torch.bfloat16,
        ).to("cuda:0").eval()
        log.info("alignment_loaded")

    def prepare(self, pcm, text):
        audio=AF.resample(torch.from_numpy(np.asarray(pcm,dtype=np.float32).copy()), SAMPLE_RATE,16000).numpy()
        # Documented default segmentation: whitespace words + individual CJK
        # characters. Avoid substituting an ASR transcript for the user's text.
        return self.processor.prepare_forced_aligner_inputs(audio=audio,transcript=text,return_tensors="pt")

    def run(self, prepared, marker, progress):
        return _align_gpu(prepared,str(marker),progress)


@spaces.GPU(duration=20)
def _align_gpu(prepared, marker, progress):
    log.info("alignment_gpu_enter request=%s",Path(marker).stem)
    inputs = outputs = None
    try:
        check_cancelled(marker)
        progress(None,desc="Synchronizing words")
        cpu_inputs,word_lists=prepared
        inputs=cpu_inputs.to("cuda:0",torch.bfloat16)
        with torch.inference_mode():
            outputs=ALIGNER.model(**inputs)
        check_cancelled(marker)
        return ALIGNER.processor.decode_forced_alignment(
            logits=outputs.logits,input_ids=inputs["input_ids"],word_lists=word_lists,
            timestamp_token_id=ALIGNER.model.config.timestamp_token_id,
        )[0]
    finally:
        del outputs, inputs
        torch.cuda.synchronize()
        log.info("alignment_gpu_exit request=%s",Path(marker).stem)


def load_aligner():
    global ALIGNER
    if ALIGNER is None:
        ALIGNER=Aligner()
    return ALIGNER