#!/usr/bin/env python3
import datetime
import hashlib
import json
import os
from pathlib import Path
import subprocess
import time
import urllib.error
import urllib.parse
import urllib.request

ROOT = Path('/data/video-pipeline/HaTramAudio/project/006-Bay-La-Thu-Khong-Nguoi-Nhan')
MANIFEST = ROOT / 'script/tts-manifest.json'
ENDPOINT = 'http://192.168.40.32:7861/voice/ngoc-huyen-clone'
INTRO_TEXT = 'Các bạn đang nghe truyện được phát từ Hạ Trâm Audio, chúc các bạn có những giây phút nghe truyện vui vẻ. Hãy ủng hộ chúng tôi bằng cách like video và đăng ký kênh.'


def now():
    return datetime.datetime.now(datetime.timezone.utc).isoformat().replace('+00:00', 'Z')


def sha256(path):
    h = hashlib.sha256()
    with path.open('rb') as f:
        for block in iter(lambda: f.read(8 * 1024 * 1024), b''):
            h.update(block)
    return h.hexdigest()


def atomic_json(path, value):
    tmp = path.with_name(f'.{path.name}.{os.getpid()}.tmp')
    with tmp.open('w', encoding='utf-8') as f:
        json.dump(value, f, ensure_ascii=False, indent=2)
        f.write('\n')
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def probe(path):
    raw = subprocess.check_output([
        'ffprobe', '-v', 'error', '-show_entries',
        'format=duration,size:stream=codec_name,codec_type,sample_rate,channels',
        '-of', 'json', str(path),
    ], text=True)
    data = json.loads(raw)
    streams = [s for s in data.get('streams', []) if s.get('codec_type') == 'audio']
    if len(streams) != 1:
        raise ValueError('expected exactly one audio stream')
    s = streams[0]
    if s.get('codec_name') != 'pcm_s16le' or s.get('sample_rate') != '48000' or s.get('channels') != 1:
        raise ValueError(f'unexpected WAV signature: {s}')
    duration = float(data['format']['duration'])
    if duration <= 0:
        raise ValueError('non-positive duration')
    return data, duration


def synthesize(text, output, attempt_log):
    body = urllib.parse.urlencode({
        'text': text, 'style': 'doc_truyen', 'speed': '1.0', 'denoise': 'true'
    }).encode()
    tmp = output.with_name(f'.{output.name}.part')
    for attempt in range(1, 4):
        started = now()
        try:
            req = urllib.request.Request(
                ENDPOINT, data=body,
                headers={'Content-Type': 'application/x-www-form-urlencoded'},
                method='POST')
            with urllib.request.urlopen(req, timeout=900) as response:
                status = response.status
                payload = response.read()
            if status != 200:
                raise RuntimeError(f'HTTP {status}')
            tmp.write_bytes(payload)
            if tmp.open('rb').read(4) != b'RIFF':
                raise ValueError('response is not RIFF WAV')
            p, duration = probe(tmp)
            os.replace(tmp, output)
            attempt_log.append({'attempt': attempt, 'status': 'completed', 'started_at': started, 'completed_at': now(), 'http': status})
            return p, duration
        except urllib.error.HTTPError as exc:
            transient = exc.code == 429 or 500 <= exc.code < 600
            attempt_log.append({'attempt': attempt, 'status': 'failed', 'started_at': started, 'completed_at': now(), 'http': exc.code, 'transient': transient})
            tmp.unlink(missing_ok=True)
            if not transient or attempt == 3:
                raise
        except (urllib.error.URLError, TimeoutError, subprocess.CalledProcessError, ValueError, RuntimeError) as exc:
            attempt_log.append({'attempt': attempt, 'status': 'failed', 'started_at': started, 'completed_at': now(), 'error': type(exc).__name__, 'transient': True})
            tmp.unlink(missing_ok=True)
            if attempt == 3:
                raise
        time.sleep(2 ** (attempt - 1))
    raise RuntimeError('unreachable')


def concat_wav(inputs, output):
    list_path = output.with_name(f'.{output.stem}-concat.txt')
    list_path.write_text(''.join(f"file '{p}'\n" for p in inputs), encoding='utf-8')
    tmp = output.with_name(f'.{output.name}.part.wav')
    subprocess.run(['ffmpeg', '-y', '-v', 'error', '-f', 'concat', '-safe', '0', '-i', str(list_path), '-c', 'copy', str(tmp)], check=True)
    p, duration = probe(tmp)
    os.replace(tmp, output)
    list_path.unlink(missing_ok=True)
    return p, duration


def main():
    manifest = json.loads(MANIFEST.read_text(encoding='utf-8'))
    manifest['status'] = 'running'
    manifest['started_at'] = manifest.get('started_at') or now()
    atomic_json(MANIFEST, manifest)

    intro = ROOT / 'audio/intro-voice.wav'
    intro_attempts = []
    if not intro.exists():
        intro_probe, intro_duration = synthesize(INTRO_TEXT, intro, intro_attempts)
    else:
        intro_probe, intro_duration = probe(intro)
    intro_receipt = {
        'status': 'completed', 'verified': True,
        'source_canon_sha256': manifest['source_canon_sha256'],
        'text': INTRO_TEXT, 'output': str(intro), 'artifact_sha256': sha256(intro),
        'duration_seconds': intro_duration, 'probe': intro_probe,
        'attempts': intro_attempts, 'completed_at': now(),
    }
    atomic_json(ROOT / 'log/intro-voice.json', intro_receipt)

    for segment in manifest['segments']:
        output = Path(segment['output'])
        text_path = Path(segment['text_path'])
        text = text_path.read_text(encoding='utf-8')
        if sha256(text_path) != segment['text_sha256']:
            raise RuntimeError(f'text hash drift for segment {segment["index"]}')
        if segment.get('status') == 'completed' and output.exists() and segment.get('sha256') == sha256(output):
            probe(output)
            continue
        attempts = []
        submitted = now()
        p, duration = synthesize(text, output, attempts)
        segment.update({
            'status': 'completed', 'submitted_at': submitted,
            'completed_at': now(), 'duration': duration,
            'size': output.stat().st_size, 'sha256': sha256(output),
            'probe': p, 'attempts': attempts,
        })
        manifest['completed_segments'] = sum(s.get('status') == 'completed' for s in manifest['segments'])
        manifest['updated_at'] = now()
        atomic_json(MANIFEST, manifest)
        print(json.dumps({'segment': segment['index'], 'completed': manifest['completed_segments'], 'total': len(manifest['segments']), 'duration': duration}), flush=True)

    inputs = [Path(s['output']) for s in manifest['segments']]
    full = ROOT / 'audio/story-full.wav'
    full_probe, full_duration = concat_wav(inputs, full)
    segment_sum = sum(float(s['duration']) for s in manifest['segments'])
    if abs(full_duration - segment_sum) > 0.1:
        raise RuntimeError(f'concat duration delta too large: {full_duration - segment_sum}')
    manifest.update({
        'status': 'completed', 'verified': True, 'completed_segments': len(inputs),
        'output': str(full), 'artifact_sha256': sha256(full),
        'duration_seconds': full_duration, 'segment_duration_sum': segment_sum,
        'duration_delta_seconds': full_duration - segment_sum,
        'probe': full_probe, 'completed_at': now(),
    })
    atomic_json(MANIFEST, manifest)
    receipt = {
        'status': 'completed', 'verified': True,
        'source_canon_sha256': manifest['source_canon_sha256'],
        'spoken_narration_sha256': manifest['spoken_narration_sha256'],
        'output': str(full), 'artifact_sha256': manifest['artifact_sha256'],
        'segment_count': len(inputs), 'duration_seconds': full_duration,
        'segment_duration_sum': segment_sum,
        'duration_delta_seconds': manifest['duration_delta_seconds'],
        'probe': full_probe, 'completed_at': manifest['completed_at'],
    }
    atomic_json(ROOT / 'log/tts.json', receipt)
    print(json.dumps({'status': 'completed', 'output': str(full), 'duration': full_duration}), flush=True)


if __name__ == '__main__':
    try:
        main()
    except Exception as exc:
        try:
            manifest = json.loads(MANIFEST.read_text(encoding='utf-8'))
            manifest.update({'status': 'failed', 'error': f'{type(exc).__name__}: {exc}', 'failed_at': now()})
            atomic_json(MANIFEST, manifest)
        finally:
            raise