#!/usr/bin/env python3
"""Chunked TTS with parallel fetching + retry + ffmpeg concat."""
import sys, urllib.parse, subprocess, os, hashlib, time, re
from concurrent.futures import ThreadPoolExecutor, as_completed

TTS_URL = "https://hf4uwho-pocket-tts.hf.space/tts"
OUTDIR = "/tmp/tts"
os.makedirs(OUTDIR, exist_ok=True)

def tts(text, voice, fmt, outfile, timeout=120, retries=2):
    """Generate TTS for a single chunk with retry."""
    encoded = urllib.parse.quote(text)
    url = f"{TTS_URL}?text={encoded}&voice={voice}&format={fmt}"
    
    for attempt in range(retries + 1):
        try:
            r = subprocess.run(
                ["curl", "-s", "-o", outfile, "-w", "%{http_code}\n%{time_total}\n%{errormsg}", 
                 "--max-time", str(timeout), "--connect-timeout", "15", url],
                capture_output=True, text=True, timeout=timeout+10
            )
            lines = r.stdout.strip().split("\n")
            code = lines[0] if lines else "?"
            curl_time = lines[1] if len(lines) > 1 else "?"
            error = lines[2] if len(lines) > 2 else ""
            
            if code == "200":
                mime = subprocess.run(["file", "-b", "--mime-type", outfile], capture_output=True, text=True).stdout.strip()
                if mime.startswith("audio/") or mime == "application/ogg":
                    return os.path.getsize(outfile)
                else:
                    raise RuntimeError(f"Not audio: {mime}")
            elif code == "000":
                if attempt < retries:
                    time.sleep(2)
                    continue
                raise RuntimeError(f"Connection failed: {error[:80]}")
            else:
                raise RuntimeError(f"HTTP {code}")
        except subprocess.TimeoutExpired:
            if attempt < retries:
                time.sleep(2)
                continue
            raise RuntimeError("Timeout")

def chunk_text(text, max_chars=400):
    """Split text into chunks at sentence boundaries."""
    sentences = re.split(r'(?<=[.!?])\s+', text)
    chunks = []
    current = ""
    for s in sentences:
        if len(current) + len(s) + 1 > max_chars and current:
            chunks.append(current.strip())
            current = s
        else:
            current = (current + " " + s).strip()
    if current.strip():
        chunks.append(current.strip())
    return chunks

def gen_chunk(i, chunk, voice, fmt, chunkdir, timeout):
    """Generate a single chunk, returns (i, path, size) or (i, path, None) on fail."""
    cfile = f"{chunkdir}/c{i:03d}.{fmt}"
    try:
        sz = tts(chunk, voice, fmt, cfile, timeout=timeout)
        return (i, cfile, sz)
    except Exception as e:
        return (i, cfile, None)

def main():
    textfile = sys.argv[1] if len(sys.argv) > 1 else None
    voice = sys.argv[2] if len(sys.argv) > 2 else "af_alloy"
    fmt = sys.argv[3] if len(sys.argv) > 3 else "ogg"
    parallel = int(sys.argv[4]) if len(sys.argv) > 4 else 5  # parallel workers

    if not textfile:
        print("Usage: voice-chunked.py <textfile> [voice] [format] [parallel]", file=sys.stderr)
        sys.exit(1)

    with open(textfile, "r") as f:
        text = f.read().strip()

    chunks = chunk_text(text)
    print(f"Split {len(text)} chars into {len(chunks)} chunks, {parallel} parallel", file=sys.stderr)

    h = hashlib.md5(f"{textfile}{voice}".encode()).hexdigest()[:12]
    chunkdir = f"{OUTDIR}/chunks_{h}"
    os.makedirs(chunkdir, exist_ok=True)

    # Estimate per-chunk timeout: 25s per 400 chars
    per_chunk_timeout = max(60, int((len(text) / max(len(chunks), 1)) * 25 / 400 * 3))
    total_start = time.time()

    # Fire all chunks in parallel via ThreadPoolExecutor
    chunk_files = {}  # i -> path
    failed = 0
    done = 0
    n = len(chunks)

    with ThreadPoolExecutor(max_workers=parallel) as pool:
        fut_map = {pool.submit(gen_chunk, i, c, voice, fmt, chunkdir, per_chunk_timeout): i 
                   for i, c in enumerate(chunks)}
        
        for fut in as_completed(fut_map):
            i, path, sz = fut.result()
            done += 1
            if sz is not None:
                chunk_files[i] = path
                print(f"  [{done}/{n}] ok {i:3d} {sz/1024:.0f}KB", file=sys.stderr, flush=True)
            else:
                failed += 1
                print(f"  [{done}/{n}] FAIL {i:3d}", file=sys.stderr, flush=True)

    if not chunk_files:
        print("ERROR: No chunks succeeded", file=sys.stderr)
        sys.exit(1)

    # Retry failed chunks serially (some may work on retry with less contention)
    if failed > 0:
        print(f"\nRetrying {failed} failed chunks serially...", file=sys.stderr)
        for i, c in enumerate(chunks):
            if i in chunk_files:
                continue
            cfile = f"{chunkdir}/c{i:03d}.{fmt}"
            try:
                sz = tts(c, voice, fmt, cfile, timeout=per_chunk_timeout*2, retries=3)
                chunk_files[i] = cfile
                print(f"  Retry ok {i:3d} {sz/1024:.0f}KB", file=sys.stderr, flush=True)
                failed -= 1
            except:
                print(f"  Retry failed {i:3d}", file=sys.stderr, flush=True)

    if failed > 0:
        print(f"WARNING: {failed}/{n} chunks permanently failed", file=sys.stderr)

    # Concatenate in order
    ordered = [chunk_files[i] for i in sorted(chunk_files)]
    outfile = f"{OUTDIR}/full_{h}.{fmt}"
    
    if len(ordered) == 1:
        import shutil
        shutil.copy2(ordered[0], outfile)
    else:
        concat_file = f"{chunkdir}/concat.txt"
        with open(concat_file, "w") as f:
            for cf in ordered:
                f.write(f"file '{cf}'\n")
        
        r = subprocess.run(
            ["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", concat_file, "-c", "copy", outfile],
            capture_output=True, text=True, timeout=30
        )
        if r.returncode != 0:
            print(f"ffmpeg failed ({r.returncode}), using raw cat", file=sys.stderr)
            with open(outfile, "wb") as out:
                for cf in ordered:
                    with open(cf, "rb") as inp:
                        out.write(inp.read())

    # Cleanup
    import shutil
    shutil.rmtree(chunkdir, ignore_errors=True)

    total = time.time() - total_start
    sz = os.path.getsize(outfile)
    print(f"Done: {outfile} ({sz/1024:.0f}KB) in {total:.1f}s", file=sys.stderr)
    print(outfile)

if __name__ == "__main__":
    main()