"""Create conservative draft Piper clips using local ASR word timestamps.
Run with .training/venv/Scripts/python.exe. No audio is uploaded.
Only sentences matching the supplied transcript exactly after normalization
are exported. Review timing and excluded sentences before training.
"""
import difflib
import json
import re
from pathlib import Path
import numpy as np
from scipy.io import wavfile
from scipy.signal import resample_poly
from faster_whisper import WhisperModel

ROOT = Path(__file__).resolve().parent / ".training" / "teddy-v1"
def tokens(text):
    return re.findall(r"[a-z0-9]+(?:'[a-z]+)?", text.lower().replace("’", "'"))

def main():
    text = (ROOT / "teddy-transcript.txt").read_text(encoding="utf-8-sig")
    cache = ROOT / "asr-words.json"
    if cache.exists():
        words = json.loads(cache.read_text(encoding="utf-8"))
    else:
        model = WhisperModel("small.en", device="cpu", compute_type="int8", cpu_threads=8,
                             download_root=str(ROOT.parent / "models"))
        rate, source = wavfile.read(ROOT / "TeddyFull.wav")
        assert source.ndim == 1 and source.dtype == np.int16
        recognition_audio = resample_poly(source.astype(np.float32) / 32768, 16000, rate)
        segments, _ = model.transcribe(recognition_audio, language="en",
                                       word_timestamps=True, vad_filter=True, beam_size=5)
        words = []
        for segment in segments:
            for word in segment.words:
                words.append({"text":word.word, "start":word.start, "end":word.end})
            print(f"Aligned through {segment.end:.1f} seconds", flush=True)
        cache.write_text(json.dumps(words, indent=2), encoding="utf-8")
    recognized = []
    for word in words:
        for token in tokens(word["text"]):
            recognized.append(dict(word, token=token))
    reference = tokens(text)
    matcher = difflib.SequenceMatcher(None, reference, [w["token"] for w in recognized], autojunk=False)
    mapping = {}
    for block in matcher.get_matching_blocks():
        for offset in range(block.size):
            mapping[block.a+offset] = block.b+offset
    rate, pcm = wavfile.read(ROOT / "TeddyFull.wav")
    assert pcm.ndim == 1 and pcm.dtype == np.int16
    audio = pcm.astype(np.float64) / 32768
    clips = ROOT / "audio"
    clips.mkdir(exist_ok=True)
    accepted, excluded, rows = [], [], []
    cursor = 0
    for index, sentence in enumerate(re.split(r"(?<=[.!?])\s+", text.strip()), 1):
        count = len(tokens(sentence))
        positions = [mapping.get(i) for i in range(cursor,cursor+count)]
        cursor += count
        if not count: continue
        if any(i is None for i in positions) or positions != list(range(positions[0],positions[0]+count)):
            excluded.append({"sentence":index,"text":sentence,"reason":"Transcript / recognition mismatch; review required"})
            continue
        first,last = positions[0],positions[-1]
        start = max(0,recognized[first]["start"]-.10)
        end = min(len(audio)/rate,recognized[last]["end"]+.10)
        # Do not include neighboring spoken words in the padding.
        if first: start=max(start,(recognized[first-1]["end"]+recognized[first]["start"])/2)
        if last+1<len(recognized): end=min(end,(recognized[last]["end"]+recognized[last+1]["start"])/2)
        if not .7 <= end-start <= 15:
            excluded.append({"sentence":index,"text":sentence,"reason":"Duration outside 0.7-15 seconds"})
            continue
        clip = resample_poly(audio[int(start*rate):int(end*rate)],22050,rate)
        filename=f"teddy_{index:04}.wav"
        wavfile.write(clips/filename,22050,np.clip(clip*32768,-32768,32767).astype(np.int16))
        rows.append(filename+"|"+sentence.replace("\n"," "))
        accepted.append({"file":filename,"start":start,"end":end,"text":sentence})
    (ROOT/"metadata.draft.csv").write_text("\n".join(rows)+"\n",encoding="utf-8")
    report={"status":"DRAFT - inspect timing and transcript before training", "matched_word_ratio":len(mapping)/len(reference),
            "source_seconds":len(audio)/rate,"source_peak":float(np.max(np.abs(audio))),
            "source_clipped_samples":int(np.sum(np.abs(audio)>=32767/32768)),
            "clip_count":len(accepted),"clip_seconds":sum(c["end"]-c["start"] for c in accepted),
            "clips":accepted,"excluded":excluded}
    (ROOT/"alignment-report.json").write_text(json.dumps(report,indent=2),encoding="utf-8")
    print(json.dumps({k:v for k,v in report.items() if k not in ("clips","excluded")},indent=2))
    print(f"Excluded {len(excluded)} sentences for review")
if __name__ == "__main__": main()