#!/usr/bin/env python3
from __future__ import annotations

import datetime
import hashlib
import json
import os
import re
import subprocess
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

ROOT = Path("/data/video-pipeline/HaTramAudio/project/013-Ngay-Buu-Cuc-Cu-Sang-Den")
FULL_PROJECT_ID = "013-Ngay-Buu-Cuc-Cu-Sang-Den"
RUN_ID = "run-20260720T110344Z-f6c08be9"
OWNER = "zoro"
ENDPOINT = "http://192.168.40.32:7861/voice/ngoc-huyen-clone"
WORK = ROOT / "work" / "tts" / RUN_ID
CHUNKS = WORK / "chunks"
SEGMENTS = WORK / "audio"
HOT = WORK / "hot-manifest.json"
FINAL_MANIFEST = ROOT / "script" / "tts-manifest.json"
NARRATION = ROOT / "story" / "spoken-narration.txt"
PROMOTION = ROOT / "script" / "promotion-report.json"
OUTPUT = ROOT / "audio" / "story-full.wav"
REQUEST_CONFIG = {"style": "doc_truyen", "speed": "1.0", "denoise": "true"}


def now() -> str:
    return datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z")


def sha256_path(path: Path) -> str:
    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 sha256_bytes(value: bytes) -> str:
    return hashlib.sha256(value).hexdigest()


def atomic_json(path: Path, value: dict) -> None:
    guard()
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name("." + path.name + ".tmp")
    with temporary.open("w", encoding="utf-8") as handle:
        json.dump(value, handle, ensure_ascii=False, indent=2)
        handle.write("\n")
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temporary, path)


def guard() -> None:
    lock = json.loads((ROOT / ".ownership-lock.json").read_text(encoding="utf-8"))
    expected = {
        "project_id": FULL_PROJECT_ID,
        "run_id": RUN_ID,
        "owner": OWNER,
        "status": "active",
    }
    for key, value in expected.items():
        if lock.get(key) != value:
            raise RuntimeError(f"ownership guard mismatch: {key}")
    if ROOT.resolve().name != FULL_PROJECT_ID:
        raise RuntimeError("project root mismatch")


def probe_wav(path: Path) -> tuple[dict, float]:
    output = subprocess.check_output([
        "ffprobe", "-v", "error", "-show_entries",
        "format=duration:stream=codec_name,sample_rate,channels",
        "-of", "json", str(path),
    ])
    data = json.loads(output)
    streams = data.get("streams", [])
    if len(streams) != 1 or streams[0].get("codec_name") != "pcm_s16le":
        raise RuntimeError("unexpected WAV codec")
    if streams[0].get("sample_rate") != "48000" or streams[0].get("channels") != 1:
        raise RuntimeError("unexpected WAV format")
    duration = float(data["format"]["duration"])
    if duration <= 0:
        raise RuntimeError("non-positive WAV duration")
    return data, duration


def chunk_bytes(raw: bytes, target_words: int = 440) -> list[bytes]:
    text = raw.decode("utf-8")
    units = re.findall(r".*?(?:\n\s*\n|\Z)", text, flags=re.DOTALL)
    units = [unit for unit in units if unit]
    if "".join(units) != text:
        raise RuntimeError("paragraph split did not reconstruct narration")
    chunks: list[str] = []
    current = ""
    current_words = 0
    for unit in units:
        words = len(re.findall(r"\b\w+\b", unit, flags=re.UNICODE))
        if current and current_words + words > target_words:
            chunks.append(current)
            current = ""
            current_words = 0
        current += unit
        current_words += words
    if current:
        chunks.append(current)
    encoded = [item.encode("utf-8") for item in chunks]
    if b"".join(encoded) != raw:
        raise RuntimeError("chunk bytes do not reconstruct narration")
    return encoded


def request_audio(text: str, output_path: Path) -> tuple[int, dict, float]:
    payload = urllib.parse.urlencode({"text": text, **REQUEST_CONFIG}).encode()
    for attempt in range(1, 4):
        guard()
        temporary = output_path.with_name(f".{output_path.name}.attempt-{attempt}.tmp.wav")
        try:
            request = urllib.request.Request(
                ENDPOINT,
                data=payload,
                headers={"Content-Type": "application/x-www-form-urlencoded"},
                method="POST",
            )
            with urllib.request.urlopen(request, timeout=900) as response:
                if response.status != 200:
                    raise RuntimeError(f"HTTP {response.status}")
                with temporary.open("wb") as handle:
                    while block := response.read(1024 * 1024):
                        handle.write(block)
                    handle.flush()
                    os.fsync(handle.fileno())
            with temporary.open("rb") as handle:
                header = handle.read(12)
            if header[:4] != b"RIFF" or header[8:12] != b"WAVE":
                raise RuntimeError("response is not RIFF/WAVE")
            probe, duration = probe_wav(temporary)
            guard()
            os.replace(temporary, output_path)
            return attempt, probe, duration
        except (urllib.error.URLError, TimeoutError, RuntimeError, subprocess.CalledProcessError):
            temporary.unlink(missing_ok=True)
            if attempt == 3:
                raise
            time.sleep(2 ** attempt)
    raise RuntimeError("unreachable")


def main() -> int:
    guard()
    promotion = json.loads(PROMOTION.read_text(encoding="utf-8"))
    if promotion.get("status") not in {"passed", "completed", "completed_content_only"} or not promotion.get("verified"):
        raise RuntimeError("promotion gate is not verified")
    raw = NARRATION.read_bytes()
    narration_hash = sha256_bytes(raw)
    source_canon_hash = promotion.get("canon_sha256") or promotion.get("source_canon_sha256")
    if narration_hash != promotion.get("spoken_narration_sha256"):
        raise RuntimeError("spoken narration hash mismatch")
    if sha256_path(ROOT / "story" / "story-canon.txt") != source_canon_hash:
        raise RuntimeError("canon hash mismatch")

    WORK.mkdir(parents=True, exist_ok=True)
    CHUNKS.mkdir(parents=True, exist_ok=True)
    SEGMENTS.mkdir(parents=True, exist_ok=True)
    parts = chunk_bytes(raw)
    plan = []
    for index, part in enumerate(parts, 1):
        text_path = CHUNKS / f"segment-{index:04d}.txt"
        if not text_path.exists() or text_path.read_bytes() != part:
            guard()
            text_path.write_bytes(part)
        plan.append({
            "index": index,
            "text_path": str(text_path),
            "text_sha256": sha256_bytes(part),
            "word_count": len(re.findall(r"\b\w+\b", part.decode("utf-8"), flags=re.UNICODE)),
            "audio_path": str(SEGMENTS / f"segment-{index:04d}.wav"),
        })
    plan_hash = sha256_bytes(json.dumps(plan, ensure_ascii=False, sort_keys=True).encode())

    hot = {
        "schema_version": 1,
        "project_id": FULL_PROJECT_ID,
        "run_id": RUN_ID,
        "status": "running",
        "verified": False,
        "candidate_sha256": promotion.get("candidate_sha256", source_canon_hash),
        "source_canon_sha256": source_canon_hash,
        "spoken_narration_sha256": narration_hash,
        "request_config": REQUEST_CONFIG,
        "chunk_count": len(plan),
        "plan_sha256": plan_hash,
        "segments": [],
        "started_at": now(),
    }
    if HOT.exists():
        previous = json.loads(HOT.read_text(encoding="utf-8"))
        if previous.get("plan_sha256") == plan_hash:
            hot = previous
            hot.update({"status": "running", "verified": False, "updated_at": now()})
    completed = {x["index"]: x for x in hot.get("segments", []) if x.get("status") == "completed"}

    try:
        for item in plan:
            index = item["index"]
            text_path = Path(item["text_path"])
            output_path = Path(item["audio_path"])
            old = completed.get(index)
            if old and old.get("text_sha256") == item["text_sha256"] and output_path.is_file():
                _, duration = probe_wav(output_path)
                if old.get("audio_sha256") == sha256_path(output_path):
                    old["duration_seconds"] = duration
                    print(f"{index:04d}/{len(plan)} resume verified", flush=True)
                    continue
            attempt_receipt = {
                "status": "requesting",
                "verified": False,
                "project_id": FULL_PROJECT_ID,
                "run_id": RUN_ID,
                "candidate_sha256": hot["candidate_sha256"],
                "source_canon_sha256": source_canon_hash,
                "spoken_narration_sha256": narration_hash,
                "chunk_index": index,
                "chunk_text_sha256": item["text_sha256"],
                "normalized_request_config": REQUEST_CONFIG,
                "created_at": now(),
            }
            attempt_path = WORK / f"attempt-{index:04d}.json"
            atomic_json(attempt_path, attempt_receipt)
            try:
                attempts, probe, duration = request_audio(text_path.read_text(encoding="utf-8"), output_path)
            except Exception as exc:
                attempt_receipt.update({
                    "status": "failed",
                    "verified": False,
                    "error_type": type(exc).__name__,
                    "error": str(exc),
                    "failed_at": now(),
                })
                atomic_json(attempt_path, attempt_receipt)
                hot["failed_segment_index"] = index
                raise
            completed[index] = {
                **item,
                "status": "completed",
                "audio_sha256": sha256_path(output_path),
                "audio_bytes": output_path.stat().st_size,
                "duration_seconds": duration,
                "probe": probe,
                "attempt_count": attempts,
                "completed_at": now(),
            }
            attempt_receipt.update({
                "status": "completed",
                "verified": True,
                "attempt_count": attempts,
                "audio_path": str(output_path),
                "audio_sha256": completed[index]["audio_sha256"],
                "audio_bytes": completed[index]["audio_bytes"],
                "duration_seconds": duration,
                "completed_at": completed[index]["completed_at"],
            })
            atomic_json(attempt_path, attempt_receipt)
            hot["segments"] = [completed[key] for key in sorted(completed)]
            hot["completed_count"] = len(completed)
            hot["updated_at"] = now()
            atomic_json(HOT, hot)
            print(f"{index:04d}/{len(plan)} completed {duration:.3f}s", flush=True)

        if sorted(completed) != list(range(1, len(plan) + 1)):
            raise RuntimeError("segment index gap")
        concat = WORK / "segments.ffconcat"
        lines = ["ffconcat version 1.0"]
        for index in sorted(completed):
            escaped = completed[index]["audio_path"].replace("'", "'\\''")
            lines.append(f"file '{escaped}'")
        guard()
        concat.write_text("\n".join(lines) + "\n", encoding="utf-8")
        temporary = OUTPUT.with_name(".story-full.tmp.wav")
        subprocess.check_call([
            "ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-f", "concat",
            "-safe", "0", "-i", str(concat), "-c", "copy", str(temporary),
        ])
        probe, duration = probe_wav(temporary)
        total = sum(completed[i]["duration_seconds"] for i in sorted(completed))
        if abs(duration - total) > 0.1:
            raise RuntimeError("concat duration mismatch")
        guard()
        os.replace(temporary, OUTPUT)
        final = {
            **hot,
            "status": "completed",
            "verified": True,
            "segment_count": len(completed),
            "segments": [completed[key] for key in sorted(completed)],
            "output": str(OUTPUT),
            "output_duration": duration,
            "segment_duration_total": total,
            "output_sha256": sha256_path(OUTPUT),
            "output_bytes": OUTPUT.stat().st_size,
            "hash_method": "sha256_stream_8MiB",
            "probe": probe,
            "completed_at": now(),
        }
        atomic_json(HOT, final)
        atomic_json(FINAL_MANIFEST, final)
        print(json.dumps({"status": "completed", "verified": True, "segments": len(completed), "duration_seconds": duration, "output_sha256": final["output_sha256"]}, ensure_ascii=False), flush=True)
        return 0
    except Exception as exc:
        failed = {
            **hot,
            "status": "failed",
            "verified": False,
            "failed_segment_index": hot.get("failed_segment_index"),
            "error_type": type(exc).__name__,
            "error": str(exc),
            "failed_at": now(),
        }
        atomic_json(HOT, failed)
        atomic_json(FINAL_MANIFEST, failed)
        raise


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