#!/usr/bin/env python3
import hashlib
import json
import re
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path

import requests

PROJECT = Path(__file__).resolve().parents[1]
STORY = PROJECT / "story" / "spoken-narration.txt"
PROMOTION = PROJECT / "story" / "promotion-report.json"
PLAN = PROJECT / "script" / "production-plan.json"
MANIFEST = PROJECT / "script" / "tts-manifest.json"
SEGMENTS = PROJECT / "audio" / "segments"
OUTPUT_WAV = PROJECT / "audio" / "story-full.wav"
OUTPUT_MP3 = PROJECT / "audio" / "story-full.mp3"
RECEIPT = PROJECT / "log" / "tts-verification.json"
ENDPOINT = "http://192.168.40.32:7861/voice/ngoc-huyen-clone"
WORD_RE = re.compile(r"\b[\wÀ-ỹĐđ]+\b", re.UNICODE)
TTS_WORDS_PER_MINUTE = 231


def sha256_bytes(data):
    return hashlib.sha256(data).hexdigest()


def sha256_file(path):
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def probe(path):
    result = subprocess.run(
        [
            "ffprobe", "-v", "error", "-show_entries",
            "format=duration,size", "-show_entries",
            "stream=index,codec_name,sample_rate,channels", "-of", "json", str(path),
        ],
        check=True,
        capture_output=True,
        text=True,
    )
    return json.loads(result.stdout)


def split_long_paragraph(paragraph, max_chars):
    sentences = re.split(r"(?<=[.!?…])\s+", paragraph.strip())
    chunks = []
    current = []
    current_len = 0
    for sentence in sentences:
        sentence = sentence.strip()
        if not sentence:
            continue
        extra = len(sentence) + (1 if current else 0)
        if current and current_len + extra > max_chars:
            chunks.append(" ".join(current))
            current = []
            current_len = 0
        if len(sentence) > max_chars:
            clauses = re.split(r"(?<=[,;:])\s+", sentence)
            for clause in clauses:
                clause = clause.strip()
                if not clause:
                    continue
                extra = len(clause) + (1 if current else 0)
                if current and current_len + extra > max_chars:
                    chunks.append(" ".join(current))
                    current = []
                    current_len = 0
                current.append(clause)
                current_len += extra
        else:
            current.append(sentence)
            current_len += extra
    if current:
        chunks.append(" ".join(current))
    return chunks


def chunk_text(text, min_chars=1700, max_chars=2400):
    paragraphs = [item.strip() for item in re.split(r"\n\s*\n", text) if item.strip()]
    units = []
    for paragraph in paragraphs:
        if len(paragraph) <= max_chars:
            units.append(paragraph)
        else:
            units.extend(split_long_paragraph(paragraph, max_chars))

    chunks = []
    current = []
    size = 0
    for unit in units:
        extra = len(unit) + (2 if current else 0)
        if current and size + extra > max_chars:
            chunks.append("\n\n".join(current))
            current = []
            size = 0
        current.append(unit)
        size += extra
        if size >= min_chars:
            chunks.append("\n\n".join(current))
            current = []
            size = 0
    if current:
        tail = "\n\n".join(current)
        if chunks and len(tail) < 700 and len(chunks[-1]) + 2 + len(tail) <= max_chars:
            chunks[-1] += "\n\n" + tail
        else:
            chunks.append(tail)
    if any(len(chunk) > max_chars for chunk in chunks):
        raise RuntimeError("Chunker exceeded max_chars after tail balancing")
    return chunks


def require_promoted_story():
    for path in (STORY, PROMOTION, PLAN):
        if not path.exists():
            raise RuntimeError(f"Required authority missing: {path}")
    promotion = json.loads(PROMOTION.read_text(encoding="utf-8"))
    if promotion.get("verified") is not True or promotion.get("status") != "passed":
        raise RuntimeError("Story is not promoted; TTS spend is blocked")
    story_bytes = STORY.read_bytes()
    canon_hash = sha256_bytes(story_bytes)
    if promotion.get("spoken_sha256") != canon_hash:
        raise RuntimeError("Story hash differs from promoted authority")
    total_words = int(promotion.get("total_words", 0))
    if not 13300 <= total_words <= 13800:
        raise RuntimeError("Promoted word count is outside target range")
    predicted_minutes = total_words / TTS_WORDS_PER_MINUTE
    if not 40 <= predicted_minutes <= 60:
        raise RuntimeError("Fixed 231 words/minute duration prediction is outside 40-60 minutes")
    plan = json.loads(PLAN.read_text(encoding="utf-8"))
    if plan.get("status") != "story_promoted_ready_for_tts" or plan.get("canon_sha256") != canon_hash:
        raise RuntimeError("Production plan is not unlocked for this canon")
    return story_bytes.decode("utf-8"), promotion, canon_hash


def build_manifest(text, promotion, canon_hash):
    chunks = chunk_text(text)
    if not chunks:
        raise RuntimeError("Chunker produced no TTS work")
    rows = []
    for index, chunk in enumerate(chunks, 1):
        rows.append({
            "index": index,
            "text": chunk,
            "text_sha256": sha256_bytes(chunk.encode("utf-8")),
            "words": len(WORD_RE.findall(chunk)),
            "output": str((SEGMENTS / f"segment-{index:04d}.wav").relative_to(PROJECT)),
            "status": "pending",
        })
    return {
        "version": 1,
        "status": "pending",
        "source": str(STORY.relative_to(PROJECT)),
        "source_canon_sha256": canon_hash,
        "source_words": promotion["total_words"],
        "endpoint": ENDPOINT,
        "request": {"style": "doc_truyen", "speed": "1.0", "denoise": "true"},
        "chunks": rows,
        "created_at": datetime.now(timezone.utc).isoformat(),
    }


def load_or_create_manifest(text, promotion, canon_hash):
    if MANIFEST.exists():
        manifest = json.loads(MANIFEST.read_text(encoding="utf-8"))
        if manifest.get("source_canon_sha256") != canon_hash:
            raise RuntimeError("Existing TTS manifest belongs to another canon hash")
        return manifest
    manifest = build_manifest(text, promotion, canon_hash)
    MANIFEST.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return manifest


def verify_existing_segment(row):
    path = PROJECT / row["output"]
    if not path.exists() or path.stat().st_size < 1024:
        return False
    try:
        data = probe(path)
        streams = data.get("streams", [])
        duration = float(data["format"]["duration"])
        valid = (
            len(streams) == 1
            and streams[0].get("codec_name") == "pcm_s16le"
            and duration > 1
        )
        if valid:
            row["duration_seconds"] = duration
            row["bytes"] = path.stat().st_size
            row["file_sha256"] = sha256_bytes(path.read_bytes())
            row["status"] = "completed"
        return valid
    except Exception:
        return False


def render_segment(row, manifest, max_attempts=3, retry_sleep=True):
    path = PROJECT / row["output"]
    path.parent.mkdir(parents=True, exist_ok=True)
    last_error = None
    for attempt in range(1, max_attempts + 1):
        try:
            response = requests.post(
                ENDPOINT,
                data={
                    "text": row["text"],
                    "style": "doc_truyen",
                    "speed": "1.0",
                    "denoise": "true",
                },
                timeout=900,
            )
            response.raise_for_status()
            if not response.content.startswith(b"RIFF"):
                raise RuntimeError("TTS response is not a RIFF/WAV file")
            temp = path.with_suffix(".wav.part")
            temp.write_bytes(response.content)
            temp.replace(path)
            if not verify_existing_segment(row):
                raise RuntimeError("Rendered WAV failed ffprobe verification")
            row["attempts"] = attempt
            return
        except Exception as exc:
            last_error = str(exc)
            row["status"] = "retrying" if attempt < max_attempts else "failed"
            row["error"] = last_error
            if attempt < max_attempts and retry_sleep:
                time.sleep(2 ** attempt)
    raise RuntimeError(f"Segment {row['index']} failed: {last_error}")


def concatenate(manifest):
    concat_file = PROJECT / "script" / "tts-concat.txt"
    lines = []
    for row in manifest["chunks"]:
        path = (PROJECT / row["output"]).resolve()
        lines.append("file '" + str(path).replace("'", "'\\''") + "'")
    concat_file.write_text("\n".join(lines) + "\n", encoding="utf-8")
    OUTPUT_WAV.parent.mkdir(parents=True, exist_ok=True)
    subprocess.run(
        ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_file), "-c", "copy", str(OUTPUT_WAV)],
        check=True,
    )
    subprocess.run(
        ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", str(OUTPUT_WAV), "-c:a", "libmp3lame", "-b:a", "192k", str(OUTPUT_MP3)],
        check=True,
    )


def verify_outputs(manifest, canon_hash):
    wav_probe = probe(OUTPUT_WAV)
    streams = wav_probe.get("streams", [])
    duration = float(wav_probe["format"]["duration"])
    segment_duration = sum(float(row["duration_seconds"]) for row in manifest["chunks"])
    verified = (
        all(row.get("status") == "completed" for row in manifest["chunks"])
        and len(streams) == 1
        and streams[0].get("codec_name") == "pcm_s16le"
        and streams[0].get("sample_rate") == "48000"
        and abs(duration - segment_duration) < 0.1
        and 40 * 60 <= duration <= 60 * 60 + 30
    )
    receipt = {
        "version": 1,
        "verified": verified,
        "source_canon_sha256": canon_hash,
        "segments_completed": sum(row.get("status") == "completed" for row in manifest["chunks"]),
        "segments_total": len(manifest["chunks"]),
        "tts_words_per_minute": TTS_WORDS_PER_MINUTE,
        "predicted_duration_minutes": round(sum(row["words"] for row in manifest["chunks"]) / TTS_WORDS_PER_MINUTE, 3),
        "duration_seconds": duration,
        "duration_minutes": round(duration / 60, 3),
        "segment_duration_sum": segment_duration,
        "duration_delta_seconds": round(duration - segment_duration, 6),
        "output_wav": str(OUTPUT_WAV.relative_to(PROJECT)),
        "output_mp3": str(OUTPUT_MP3.relative_to(PROJECT)),
        "wav_sha256": sha256_file(OUTPUT_WAV),
        "probe": wav_probe,
        "checked_at": datetime.now(timezone.utc).isoformat(),
    }
    RECEIPT.parent.mkdir(parents=True, exist_ok=True)
    RECEIPT.write_text(json.dumps(receipt, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    if not verified:
        raise RuntimeError("Final TTS verification failed")
    return receipt


def main():
    try:
        text, promotion, canon_hash = require_promoted_story()
        SEGMENTS.mkdir(parents=True, exist_ok=True)
        manifest = load_or_create_manifest(text, promotion, canon_hash)
        manifest["status"] = "running"
        MANIFEST.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        pending = []
        for row in manifest["chunks"]:
            if row.get("status") == "completed" and verify_existing_segment(row):
                continue
            pending.append(row)
        with ThreadPoolExecutor(max_workers=3) as pool:
            futures = {pool.submit(render_segment, row, manifest): row for row in pending}
            for future in as_completed(futures):
                row = futures[future]
                future.result()
                MANIFEST.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
                print(json.dumps({"segment": row["index"], "status": row["status"], "duration": row.get("duration_seconds")}, ensure_ascii=False), flush=True)
        concatenate(manifest)
        receipt = verify_outputs(manifest, canon_hash)
        manifest["status"] = "completed"
        manifest["completed_at"] = datetime.now(timezone.utc).isoformat()
        manifest["output_wav"] = str(OUTPUT_WAV.relative_to(PROJECT))
        manifest["output_mp3"] = str(OUTPUT_MP3.relative_to(PROJECT))
        MANIFEST.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        print(json.dumps({"status": "completed", "verified": receipt["verified"], "duration_seconds": receipt["duration_seconds"]}, ensure_ascii=False))
        return 0
    except Exception as exc:
        print(f"TTS blocked: {exc}", file=sys.stderr)
        return 1


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