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

import base64
import datetime
import hashlib
import json
import os
import re
import subprocess
import sys
import tempfile
import time
import urllib.parse
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
PROJECT_MANIFEST = ROOT / "script/project-manifest.json"
ENDPOINT = "http://192.168.40.32:7861/voice/ngoc-huyen-clone"
MAX_CHARS = 260


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


def resolve_active(manifest: dict, key: str) -> Path:
    value = (manifest.get("active_paths") or {}).get(key)
    if not isinstance(value, str) or not value:
        raise RuntimeError(f"active_paths.{key} is missing")
    path = Path(value)
    if not path.is_absolute():
        path = ROOT / path
    try:
        path.resolve().relative_to(ROOT.resolve())
    except ValueError as exc:
        raise RuntimeError(f"active path escapes project: {key}") from exc
    return path


def atomic_json(path: Path, value: dict) -> None:
    data = (json.dumps(value, ensure_ascii=False, indent=2) + "\n").encode("utf-8")
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    try:
        with os.fdopen(fd, "wb") as handle:
            handle.write(data)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(name, path)
    finally:
        if os.path.exists(name):
            os.unlink(name)


def split_units(text: str) -> list[str]:
    paragraphs = [value.strip() for value in text.split("\n\n") if value.strip()]
    units: list[str] = []
    for paragraph in paragraphs:
        sentences = [value.strip() for value in re.split(r"(?<=[.!?…])\s+", paragraph) if value.strip()]
        pieces: list[str] = []
        for sentence in sentences:
            remaining = sentence
            while len(remaining) > MAX_CHARS:
                window = remaining[:MAX_CHARS + 1]
                candidates = [match.end() for match in re.finditer(r"[,;:]\s+", window)]
                cut = candidates[-1] if candidates else window.rfind(" ") + 1
                if cut <= 0:
                    raise RuntimeError("unsplittable narration token exceeds max chars")
                pieces.append(remaining[:cut].strip())
                remaining = remaining[cut:].strip()
            if remaining:
                pieces.append(remaining)
        buffer = ""
        for piece in pieces:
            if buffer and len(buffer) + 1 + len(piece) > MAX_CHARS:
                units.append(buffer)
                buffer = piece
            else:
                buffer = f"{buffer} {piece}".strip()
        if buffer:
            units.append(buffer)
    packed: list[str] = []
    buffer = ""
    for unit in units:
        if buffer and len(buffer) + 1 + len(unit) > MAX_CHARS:
            packed.append(buffer)
            buffer = unit
        else:
            buffer = f"{buffer} {unit}".strip()
    if buffer:
        packed.append(buffer)
    if " ".join(packed) != " ".join(text.split()):
        raise RuntimeError("segmentation changed narration tokens")
    return packed


def find_audio(value: object) -> bytes | None:
    if isinstance(value, dict):
        for key in ("audio_base64", "base64", "audio", "data"):
            item = value.get(key)
            if isinstance(item, str):
                raw = item.split(",", 1)[-1] if item.startswith("data:") else item
                try:
                    decoded = base64.b64decode(raw, validate=True)
                except Exception:
                    decoded = b""
                if decoded.startswith(b"RIFF") and decoded[8:12] == b"WAVE":
                    return decoded
        for key in ("audio_path", "path", "output_path", "file"):
            item = value.get(key)
            if isinstance(item, str):
                path = Path(item)
                if path.exists():
                    data = path.read_bytes()
                    if data.startswith(b"RIFF") and data[8:12] == b"WAVE":
                        return data
        for key in ("audio_url", "url"):
            item = value.get(key)
            if isinstance(item, str) and item.startswith(("http://", "https://")):
                with urllib.request.urlopen(item, timeout=300) as response:
                    data = response.read()
                if data.startswith(b"RIFF") and data[8:12] == b"WAVE":
                    return data
        for item in value.values():
            found = find_audio(item)
            if found:
                return found
    elif isinstance(value, list):
        for item in value:
            found = find_audio(item)
            if found:
                return found
    return None


def probe(path: Path) -> dict:
    return json.loads(subprocess.check_output([
        "ffprobe", "-v", "error", "-show_entries",
        "format=duration,format_name:stream=codec_name,codec_type,sample_rate,channels",
        "-of", "json", str(path),
    ], text=True))


def valid_wav(path: Path) -> bool:
    try:
        info = probe(path)
        streams = [item for item in info.get("streams", []) if item.get("codec_type") == "audio"]
        return path.stat().st_size > 44 and len(streams) == 1 and float(info["format"]["duration"]) > 0
    except Exception:
        return False


def render(text: str, output: Path) -> tuple[int, str]:
    payload = urllib.parse.urlencode({
        "text": text,
        "style": "doc_truyen",
        "speed": "1.0",
        "denoise": "true",
    }).encode("utf-8")
    req = urllib.request.Request(
        ENDPOINT,
        data=payload,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        method="POST",
    )
    temp = output.with_name(f".{output.name}.response.tmp")
    temp.unlink(missing_ok=True)
    with urllib.request.urlopen(req, timeout=600) as response:
        status = response.status
        content_type = response.headers.get("Content-Type", "")
        first = response.read(12)
        if first.startswith(b"RIFF") and first[8:12] == b"WAVE":
            with temp.open("wb") as handle:
                handle.write(first)
                for chunk in iter(lambda: response.read(1024 * 1024), b""):
                    handle.write(chunk)
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(temp, output)
        else:
            body = first + response.read()
            payload_value = json.loads(body.decode("utf-8"))
            audio = find_audio(payload_value)
            if audio is None:
                raise RuntimeError("provider response contains no WAV audio")
            with temp.open("wb") as handle:
                handle.write(audio)
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(temp, output)
    if not valid_wav(output):
        output.unlink(missing_ok=True)
        raise RuntimeError("provider WAV failed independent probe")
    return status, content_type


def main() -> int:
    project = json.loads(PROJECT_MANIFEST.read_text(encoding="utf-8"))
    narration = resolve_active(project, "spoken_narration")
    promotion_path = resolve_active(project, "promotion_receipt")
    tts_manifest = resolve_active(project, "tts_manifest")
    partial_path = resolve_active(project, "tts_partial_manifest")
    receipt_path = resolve_active(project, "tts_receipt")
    full_path = resolve_active(project, "tts_audio")
    concat_path = resolve_active(project, "tts_concat")
    audio_dir = full_path.parent
    promotion = json.loads(promotion_path.read_text(encoding="utf-8"))
    narration_hash = sha256(narration)
    if promotion.get("status") != "completed" or promotion.get("verified") is not True or promotion.get("spoken_narration_sha256") != narration_hash:
        raise RuntimeError("promotion absent, failed, or stale")
    if receipt_path.exists() or tts_manifest.exists() or full_path.exists():
        raise RuntimeError("TTS production output already exists and requires audit/invalidation")

    audio_dir.mkdir(parents=True, exist_ok=True)
    text = narration.read_text(encoding="utf-8")
    units = split_units(text)
    existing: dict[int, dict] = {}
    if partial_path.exists():
        partial = json.loads(partial_path.read_text(encoding="utf-8"))
        if partial.get("spoken_narration_sha256") == narration_hash:
            existing = {int(item["index"]): item for item in partial.get("segments", [])}

    segments: list[dict] = []
    for index, unit in enumerate(units, start=1):
        output = audio_dir / f"segment-{index:04d}.wav"
        text_hash = hashlib.sha256(unit.encode("utf-8")).hexdigest()
        prior = existing.get(index)
        reusable = (
            prior is not None
            and prior.get("text_sha256") == text_hash
            and prior.get("artifact_sha256") == (sha256(output) if output.exists() else None)
            and valid_wav(output)
        )
        http_status = prior.get("http_status") if reusable else None
        content_type = prior.get("response_content_type") if reusable else None
        if not reusable:
            output.unlink(missing_ok=True)
            try:
                http_status, content_type = render(unit, output)
            except Exception as exc:
                output.unlink(missing_ok=True)
                raise RuntimeError(f"segment {index} failed: {exc}") from exc
        info = probe(output)
        duration = float(info["format"]["duration"])
        item = {
            "index": index,
            "text_sha256": text_hash,
            "chars": len(unit),
            "words": len(unit.split()),
            "path": str(output),
            "artifact_sha256": sha256(output),
            "duration_seconds": duration,
            "http_status": http_status,
            "response_content_type": content_type,
            "verified": True,
        }
        segments.append(item)
        atomic_json(partial_path, {
            "status": "running",
            "verified": False,
            "spoken_narration_sha256": narration_hash,
            "planned_count": len(units),
            "completed_count": len(segments),
            "segments": segments,
            "updated_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        })

    if [item["index"] for item in segments] != list(range(1, len(units) + 1)):
        raise RuntimeError("segment index gap or duplicate")
    concat_path.write_text("".join(f"file '{Path(item['path']).name}'\n" for item in segments), encoding="utf-8")
    temp_full = audio_dir / ".story-full.tmp.wav"
    temp_full.unlink(missing_ok=True)
    subprocess.run([
        "ffmpeg", "-v", "error", "-y", "-f", "concat", "-safe", "0",
        "-i", str(concat_path), "-c:a", "pcm_s16le", str(temp_full),
    ], check=True)
    if not valid_wav(temp_full):
        raise RuntimeError("concatenated story-full WAV failed probe")
    os.replace(temp_full, full_path)
    full_probe = probe(full_path)
    full_duration = float(full_probe["format"]["duration"])
    segment_duration = sum(item["duration_seconds"] for item in segments)
    duration_delta = abs(full_duration - segment_duration)
    verified = duration_delta <= max(0.1, len(segments) * 0.002)
    manifest = {
        "status": "completed" if verified else "failed",
        "verified": verified,
        "project_id": project.get("project_id"),
        "source_canon_sha256": promotion["canon_sha256"],
        "spoken_narration_sha256": narration_hash,
        "voice": "Ngoc-Huyen-Clone",
        "endpoint": ENDPOINT,
        "style": "doc_truyen",
        "speed": 1.0,
        "denoise": True,
        "planned_count": len(units),
        "completed_count": len(segments),
        "segments": segments,
        "story_full_path": str(full_path),
        "story_full_sha256": sha256(full_path),
        "duration_seconds": full_duration,
        "segment_duration_sum": segment_duration,
        "duration_delta_seconds": duration_delta,
        "probe": full_probe,
        "completed_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "hash_method": "sha256 streaming 8 MiB",
    }
    atomic_json(tts_manifest, manifest)
    receipt = {
        "status": manifest["status"],
        "verified": verified,
        "source_canon_sha256": promotion["canon_sha256"],
        "spoken_narration_sha256": narration_hash,
        "input_path": str(narration),
        "output_path": str(full_path),
        "artifact_sha256": manifest["story_full_sha256"],
        "segment_count": len(segments),
        "duration_seconds": full_duration,
        "endpoint_verification": {"url": ENDPOINT, "direct_http": True},
        "independent_probe": full_probe,
        "manifest_path": str(tts_manifest),
        "completed_at": manifest["completed_at"],
        "hash_method": manifest["hash_method"],
    }
    atomic_json(receipt_path, receipt)
    print(json.dumps({
        "status": receipt["status"], "verified": verified,
        "segments": len(segments), "duration_seconds": full_duration,
        "artifact_sha256": receipt["artifact_sha256"], "receipt": str(receipt_path),
    }, ensure_ascii=False))
    return 0 if verified else 1


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Exception as exc:
        failure = {
            "status": "failed", "verified": False,
            "error": f"{type(exc).__name__}: {exc}",
            "failed_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        }
        try:
            project = json.loads(PROJECT_MANIFEST.read_text(encoding="utf-8"))
            atomic_json(resolve_active(project, "tts_receipt"), failure)
        except Exception:
            pass
        print(json.dumps(failure, ensure_ascii=False), file=sys.stderr)
        sys.exit(1)