#!/usr/bin/env python3
"""Test-only replacements for op/curl; no network or real credentials."""

import base64
import json
import os
from pathlib import Path
import subprocess
import sys
import time


if Path(sys.argv[0]).name == "op":
    print("test-api-key")
    sys.exit(0)

args = sys.argv[1:]
request_path = Path(args[args.index("--data-binary") + 1][1:])
request = json.loads(request_path.read_text())
output = Path(args[args.index("-o") + 1])
mode = os.environ.get("TRANSCRIBE_TEST_MODE", "ok")
events_path = Path(os.environ["TRANSCRIBE_TEST_LOG"]).with_suffix(".events.jsonl")
is_comparison = "https://openrouter.ai/api/v1/chat/completions" in args
index = -1 if is_comparison else int(output.stem.split("-")[1])
staging = output.parent.parent if is_comparison else output.parent


def event(kind, **extra):
    record = dict(event=kind, index=index, pid=os.getpid(), group=os.getpgrp(),
                  time=time.monotonic(), staging=str(staging), **extra)
    with events_path.open("a") as log:
        log.write(json.dumps(record) + "\n")


event("start")
if is_comparison:
    samples = json.loads((output.parent / "samples.json").read_text())
    content = request["messages"][-1]["content"]
    audio_parts = [part["input_audio"] for part in content if part["type"] == "input_audio"]
    assert len(content) == 2 * len(samples) == 2 * len(audio_parts)
    for number, sample in enumerate(samples):
        assert content[number * 2] == {"type": "text", "text": f"Sample ID {sample['id']}"}
        assert content[number * 2 + 1]["type"] == "input_audio"
    assert request["response_format"]["type"] == "json_schema"
    assert request["response_format"]["json_schema"]["strict"] is True
    prompt = "\n".join(part["text"] for message in request["messages"]
                       for part in (message["content"] if isinstance(message["content"], list)
                                    else [{"type": "text", "text": message["content"]}])
                       if part["type"] == "text")
    assert not any(word in prompt for word in ("Alpha.", "Again.", "Beta.", "AlphaLate.", "BetaLate."))
    durations = []
    for number, part in enumerate(audio_parts):
        assert part["format"] == "mp3"
        audio = output.parent / f"checked-{number}.mp3"
        audio.write_bytes(base64.b64decode(part["data"], validate=True))
        info = json.loads(subprocess.check_output([
            "ffprobe", "-v", "error", "-show_entries",
            "format=duration:stream=codec_name,sample_rate,channels", "-of", "json", str(audio),
        ], text=True))
        assert info["streams"][0] == {"codec_name": "mp3", "sample_rate": "16000", "channels": 1}
        durations.append(float(info["format"]["duration"]))
        assert 1.5 <= durations[-1] <= 8.1
    Path(os.environ["TRANSCRIBE_TEST_LOG"]).with_suffix(".speakers.json").write_text(json.dumps({
        "request": request, "samples": samples, "durations": durations, "staging": str(staging),
    }))
    if mode == "reconcile-blocked":
        child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
        event("child", child_pid=child.pid)
        time.sleep(30)
    if mode == "reconcile-http-error":
        output.write_text('{"error":{"message":"simulated comparison failure"}}')
        sys.exit(22)
    if mode == "reconcile-bad-json":
        output.write_text("not json")
        sys.exit(0)
    assignments = [{"sample": sample["id"], "speaker": 0 if round(sample["start"]) in (1, 6, 60) else 1}
                   for sample in samples]
    if mode == "reconcile-missing":
        assignments.pop()
    elif mode == "reconcile-duplicate":
        assignments.append(assignments[0])
    elif mode == "reconcile-unknown-sample":
        assignments[0]["sample"] = "not-a-sample"
    elif mode == "reconcile-invalid-speaker":
        assignments[0]["speaker"] = -1
    elif mode == "reconcile-null":
        assignments[0]["speaker"] = None
    output.write_text(json.dumps({"choices": [{
        "finish_reason": "length" if mode == "reconcile-truncated" else "stop",
        "message": {"content": json.dumps({"assignments": assignments})},
    }]}))
    event("finish")
    sys.exit(0)

manifest = json.loads((output.parent / "chunks.json").read_text())
chunk = manifest[index]
if mode in ("parallel-error", "blocked") and index == 0:
    child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
    event("child", child_pid=child.pid)
assert request["model"] == "microsoft/mai-transcribe-2"
assert request["input_audio"]["format"] == "mp3"
if len(manifest) > 1:
    assert request["response_format"] == "verbose_json"
    assert request["timestamp_granularities"] == ["segment", "word"]

audio = output.parent / f"uploaded-{index}.mp3"
audio.write_bytes(base64.b64decode(request["input_audio"]["data"], validate=True))
info = json.loads(subprocess.check_output([
    "ffprobe", "-v", "error", "-show_entries",
    "format=duration:stream=codec_name,sample_rate,channels", "-of", "json", str(audio),
], text=True))
assert info["streams"][0] == {"codec_name": "mp3", "sample_rate": "16000", "channels": 1}
assert abs(float(info["format"]["duration"]) - (chunk["end"] - chunk["start"])) < 0.1
with open(os.environ["TRANSCRIBE_TEST_LOG"], "a") as log:
    log.write(json.dumps({
        "index": index, "chunk": chunk, "staging": str(output.parent),
        "work": str(request_path.parent),
        "upload_duration": float(info["format"]["duration"]),
        "provider": request.get("provider"),
        "response_format": request.get("response_format"),
        "timestamp_granularities": request.get("timestamp_granularities"),
    }) + "\n")

if mode in ("parallel", "parallel-error"):
    deadline = time.monotonic() + 5
    barrier = int(os.environ.get("TRANSCRIBE_TEST_BARRIER", "1"))
    while True:
        events = [json.loads(line) for line in events_path.read_text().splitlines()]
        enough = sum(e["event"] == "start" for e in events) >= barrier
        child_ready = mode != "parallel-error" or any(e["event"] == "child" for e in events)
        if enough and child_ready:
            break
        if time.monotonic() >= deadline:
            raise RuntimeError("Worker concurrency barrier timed out")
        time.sleep(0.01)
    if mode == "parallel-error" and index == 1:
        output.write_text('{"error":{"message":"simulated concurrent failure"}}')
        event("finish")
        sys.exit(22)
    time.sleep(30 if mode == "parallel-error" else (0.8 if index == 0 else 0.05))
elif mode == "blocked":
    time.sleep(30)

if index == 1:
    if mode == "http-error":
        output.write_text('{"error":{"message":"simulated timeout"}}')
        sys.exit(22)
    if mode == "bad-json":
        output.write_text("not json")
        sys.exit(0)
    if mode == "api-error":
        output.write_text('{"error":{"message":"simulated API error"},"text":""}')
        sys.exit(0)
    if mode == "missing-words":
        output.write_text('{"text":"speech without timestamps"}')
        sys.exit(0)

if "response_format" not in request:
    response = {"text": "Provider's original formatting.\nNext line."}
else:
    words = [
        {"word": text, "start": start - chunk["start"], "end": end - chunk["start"]}
        for text, start, end in [
            ("Hello.", 1, 1.4), ("boundary", 56.7, 57.1),
            ("continues.", 57.4, 57.8), ("Last.", 65, 65.4),
        ] if start >= chunk["start"] and end <= chunk["end"]
    ]
    if mode in ("speakers", "partial-speakers"):
        if len(manifest) == 1:
            words = [
                {"word": "Hello", "start": 1, "end": 1.4},
                {"word": "there", "start": 1.5, "end": 1.9},
            ]
        if request["provider"]["options"]["azure"]["diarization"]["enabled"]:
            for speaker, word in enumerate(words):
                if mode == "speakers" or speaker == 0:
                    word["speaker"] = (speaker if len(manifest) == 1 else
                                       int(word["start"] + chunk["start"] >= 57.3))
    if mode == "parallel":
        middle = (chunk["end"] - chunk["start"]) / 2
        words = [{"word": f"Part{index}.", "start": middle, "end": middle + 0.4}]
    if mode.startswith("reconcile"):
        # Two labels already describe voice A inside the first request; later
        # chunks reuse the same local numbers for different request identities.
        words = [
            {"word": text, "start": start - chunk["start"], "end": start + 2 - chunk["start"],
             "speaker": speaker}
            for text, start, speaker in [
                ("Alpha.", 1, 7), ("Again.", 6, 8), ("Beta.", 12, 9),
                ("AlphaLate.", 60, 9), ("BetaLate.", 65, 7),
            ] if start >= chunk["start"] and start + 2 <= chunk["end"]
        ]
    response = {"text": " ".join(w["word"] for w in words), "words": words}
output.write_text(json.dumps(response))
event("finish")