"""Run whole-recording Pyannote diarization and align its speakers to MAI words.

Only speaker labels change. A word needs a unique speaker covering at least
half of its duration. If there is no overlap, a nearest turn is usable only
when every point of the word lies within 200 ms of that turn. Equal candidates
remain unknown, including zero-duration words exactly at a speaker boundary.
"""

import argparse
from bisect import bisect_left, bisect_right
import json
import math
import os
from pathlib import Path
import re
import subprocess
import sys
import tempfile
import time
from urllib.parse import quote, urlsplit
import uuid
import wave

from subtitles import convert, normalized_words
from transcript import render_text

API_URL = "https://api.pyannote.ai/v1"
DEFAULT_TIMEOUT = 3600
POLL_INTERVAL = 2.0
MAX_NEAREST_GAP = 0.2
MAX_AUDIO_SECONDS = 24 * 60 * 60
MAX_UPLOAD_BYTES = 1024 ** 3
TRANSIENT_HTTP = {408, 425, 429, 500, 502, 503, 504}
EPSILON = 1e-9


def check_speakers(speakers):
    if speakers is not None and (type(speakers) is not int or speakers < 1):
        raise ValueError("The known speaker count must be a positive integer")


def unique_object(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError("Duplicate JSON key")
        result[key] = value
    return result


def read_json(path):
    with open(path, encoding="utf-8") as source:
        return json.load(source, object_pairs_hook=unique_object)


def finite_number(value):
    return type(value) in (int, float) and math.isfinite(value)


def validate_segments(segments, exclusive=True):
    if not isinstance(segments, list):
        raise ValueError("Pyannote result is missing the exclusive diarization timeline"
                         if exclusive else "Invalid Pyannote standard diarization timeline")
    previous_end = 0
    for segment in segments:
        if not isinstance(segment, dict):
            raise ValueError("Invalid Pyannote diarization turn")
        start, end, speaker = (segment.get(key) for key in ("start", "end", "speaker"))
        if (not finite_number(start) or not finite_number(end)
                or not 0 <= start < end or (exclusive and start < previous_end)
                or not isinstance(speaker, str) or not speaker.strip()):
            raise ValueError("Pyannote turns must have finite valid times and speaker IDs; "
                             "exclusive turns must be ordered and nonoverlapping")
        confidence = segment.get("confidence")
        if confidence is not None and (
                not isinstance(confidence, dict)
                or any(not isinstance(key, str) or not key.strip()
                       or not finite_number(value) or not 0 <= value <= 100
                       for key, value in confidence.items())):
            raise ValueError("Invalid Pyannote turn confidence scores")
        previous_end = end
    return segments


def validated_timeline(job):
    """Require a completed exclusive timeline; never sort or repair its turns."""
    if (not isinstance(job, dict) or job.get("status") != "succeeded"
            or job.get("error") or not isinstance(job.get("output"), dict)
            or job["output"].get("error")):
        raise ValueError("Pyannote result must be a succeeded job with complete output")
    return validate_segments(job["output"].get("exclusiveDiarization"))


def retained_result(job, job_id):
    """Keep only validated diarization data, never input URLs or request metadata."""
    validated_timeline(job)
    output = {}
    for name in ("exclusiveDiarization", "diarization"):
        if name in job["output"]:
            segments = validate_segments(job["output"][name], exclusive=name == "exclusiveDiarization")
            output[name] = [{key: value for key, value in segment.items()
                             if key in ("start", "end", "speaker", "confidence")}
                            for segment in segments]
    return {"jobId": job_id, "status": "succeeded", "output": output}


def align(response, job, speakers=None):
    """Replace local labels using the complete audio timeline, without editing words."""
    check_speakers(speakers)
    # Validate the originals without taking normalized copies: those normalize
    # whitespace and clamp regressing STT timestamps, which must survive here.
    normalized_words(response, "Speaker alignment")
    segments = validated_timeline(job)
    if not segments and any(word["word"].strip() for word in response["words"]):
        raise ValueError("Pyannote returned an empty speaker timeline for a nonempty transcript")
    recording_speakers = {segment["speaker"] for segment in segments}
    labels = {}
    starts = [segment["start"] for segment in segments]
    ends = [segment["end"] for segment in segments]
    counts = {"unknown": 0, "nearest": 0}
    result_words = []
    for original in response["words"]:
        start, end = original["start"], original["end"]
        support = {}
        if end > start:
            index = bisect_right(ends, start)
            while index < len(segments) and starts[index] < end:
                segment = segments[index]
                overlap = min(end, segment["end"]) - max(start, segment["start"])
                if overlap > 0:
                    key = segment["speaker"]
                    support[key] = support.get(key, 0) + overlap
                index += 1
        else:
            index = bisect_left(ends, start)
            while index < len(segments) and starts[index] <= start:
                support[segments[index]["speaker"]] = 1
                index += 1
        selected = None
        used_nearest = False
        if support:
            best = max(support.values())
            winners = [key for key, amount in support.items() if abs(amount - best) <= EPSILON]
            if len(winners) == 1 and best + EPSILON >= (end - start) / 2:
                selected = winners[0]
        elif segments:
            # Only adjacent turns can be nearest in an ordered exclusive
            # timeline. The bound applies to the whole word, not just its edge.
            index = bisect_right(starts, start)
            distances = {}
            for candidate in (index - 1, index):
                if 0 <= candidate < len(segments):
                    segment = segments[candidate]
                    distance = max(segment["start"] - start, end - segment["end"], 0)
                    key = segment["speaker"]
                    distances[key] = min(distances.get(key, math.inf), distance)
            best = min(distances.values())
            winners = [key for key, distance in distances.items() if abs(distance - best) <= EPSILON]
            if len(winners) == 1 and best <= MAX_NEAREST_GAP + EPSILON:
                selected = winners[0]
                used_nearest = True
        if not original["word"].strip():
            selected = None
        if selected is not None:
            labels.setdefault(selected, len(labels))
            counts["nearest"] += used_nearest
        counts["unknown"] += selected is None
        result_words.append(dict(original, speaker=labels[selected] if selected is not None else None))
    print(f"Speaker alignment: {len(recording_speakers)} recording speaker(s), "
          f"{len(result_words) - counts['unknown']}/{len(result_words)} words assigned.", file=sys.stderr)
    if counts["nearest"]:
        print(f"Warning: {counts['nearest']} word(s) used a nearest turn within "
              f"{MAX_NEAREST_GAP:.1f} seconds of the entire word.", file=sys.stderr)
    if counts["unknown"]:
        print(f"Warning: {counts['unknown']} word(s) have ambiguous or missing speaker timing; "
              "their speaker is unknown.", file=sys.stderr)
    if speakers is not None and len(recording_speakers) != speakers:
        print(f"Warning: expected {speakers} speakers, but Pyannote returned {len(recording_speakers)}; "
              "the count is not enforced by merging identities.", file=sys.stderr)
    return dict(response, words=result_words)


class RequestError(ValueError):
    def __init__(self, message, transient=False):
        super().__init__(message)
        self.transient = transient


def safe_detail(payload, api_key):
    """Keep useful provider errors without echoing credentials or signed URLs."""
    if not isinstance(payload, dict):
        return ""
    detail = payload.get("error", payload.get("detail", payload.get("message", "")))
    if isinstance(detail, dict):
        detail = detail.get("message", detail.get("detail", ""))
    if not isinstance(detail, str):
        return ""
    detail = detail.replace(api_key, "[redacted]") if api_key else detail
    detail = re.sub(r"(?:https?://|media://)\S+", "[redacted URL]", detail)
    detail = re.sub(r"(?i)bearer\s+\S+", "Bearer [redacted]", detail)
    return " ".join(detail.split())[:400]


class Client:
    def __init__(self, directory, api_key, deadline):
        self.directory = Path(directory)
        self.api_key = api_key
        self.deadline = deadline
        self.auth_path = self.directory / "authorization.header"

    def remaining(self):
        remaining = self.deadline - time.monotonic()
        if remaining <= 0:
            raise ValueError("Pyannote operation timed out; a submitted remote job may continue")
        return remaining

    def request(self, method, url, *, payload=None, upload=None, authenticated=True):
        parsed = urlsplit(url)
        if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password:
            raise ValueError("Pyannote returned an invalid HTTPS upload URL")
        response_path = self.directory / "http-response.json"
        response_path.unlink(missing_ok=True)
        timeout = min(self.remaining(), 300 if upload else 120)
        command = ["curl", "--silent", "--show-error", "--request", method,
                   "--connect-timeout", str(min(20, timeout)), "--max-time", str(timeout),
                   "--proto", "=https", "-o", str(response_path), "--write-out", "%{http_code}"]
        if authenticated:
            # Keep credentials out of process arguments. The parent directory
            # is private, and this file is removed with it after the operation.
            descriptor = os.open(self.auth_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
            with os.fdopen(descriptor, "w", encoding="utf-8") as header:
                header.write(f"Authorization: Bearer {self.api_key}\n")
            command.extend(["--header", f"@{self.auth_path}"])
        if payload is not None:
            request_path = self.directory / "http-request.json"
            request_path.write_text(json.dumps(payload), encoding="utf-8")
            command.extend(["--header", "Content-Type: application/json",
                            "--data-binary", f"@{request_path}"])
        if upload is not None:
            command.extend(["--header", "Content-Type: application/octet-stream",
                            "--upload-file", str(upload)])
        command.append(url)
        try:
            result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                    timeout=timeout + 1, check=False)
        except subprocess.TimeoutExpired as error:
            raise RequestError("Pyannote HTTP request timed out", transient=True) from error
        if result.returncode:
            # curl diagnostics can contain a presigned storage URL. Report the
            # exit status without copying their potentially sensitive contents.
            raise RequestError(f"Pyannote HTTP transport failed (curl exit {result.returncode})",
                               transient=True)
        status_text = result.stdout.decode("ascii", errors="replace").strip()
        if not re.fullmatch(r"[1-5][0-9]{2}", status_text):
            raise RequestError("Pyannote HTTP request returned no valid status code")
        status = int(status_text)
        if not 200 <= status < 300:
            try:
                body = read_json(response_path)
            except (ValueError, OSError):
                body = None
            detail = safe_detail(body, self.api_key)
            raise RequestError(f"Pyannote HTTP {status}" + (f": {detail}" if detail else ""),
                               transient=status in TRANSIENT_HTTP)
        if upload is not None:
            return None
        try:
            body = read_json(response_path)
        except (ValueError, OSError) as error:
            raise RequestError("Pyannote returned invalid JSON") from error
        if not isinstance(body, dict):
            raise RequestError("Pyannote returned a JSON value instead of an object")
        if body.get("error") and body.get("status") not in ("failed", "canceled"):
            detail = safe_detail(body, self.api_key)
            raise RequestError("Pyannote returned an error" + (f": {detail}" if detail else ""))
        return body

    def pause(self):
        time.sleep(min(POLL_INTERVAL, self.remaining()))


def run(audio, work, model="precision-2", speakers=None, timeout=DEFAULT_TIMEOUT):
    """Upload the complete lossless recording once, submit once, and poll its job."""
    check_speakers(speakers)
    if model not in ("precision-2", "community-1"):
        raise ValueError("The speaker model must be precision-2 or community-1")
    if not finite_number(timeout) or timeout <= 0:
        raise ValueError("The Pyannote timeout must be a positive finite number")
    api_key = os.environ.get("PYANNOTE_API_KEY", "")
    if not api_key.strip() or "\r" in api_key or "\n" in api_key:
        raise ValueError("PYANNOTE_API_KEY must contain a Pyannote API key")
    work = Path(work)
    work.mkdir(parents=True, exist_ok=True)
    result_path = work / "result.json"
    result_path.unlink(missing_ok=True)
    with wave.open(str(audio), "rb") as recording:
        if (recording.getnchannels(), recording.getsampwidth(), recording.getframerate()) != (1, 2, 16000):
            raise ValueError("Pyannote requires the prepared mono 16 kHz 16-bit PCM WAV")
        duration = recording.getnframes() / recording.getframerate()
    if not 0 < duration <= MAX_AUDIO_SECONDS:
        raise ValueError("Pyannote requires nonempty audio no longer than 24 hours")
    deadline = time.monotonic() + timeout
    # API requests/responses can contain temporary signed storage URLs. Only a
    # validated final job survives this private upload directory's lifetime.
    with tempfile.TemporaryDirectory(prefix=".upload-", dir=work) as temporary:
        client = Client(temporary, api_key, deadline)
        flac = Path(temporary) / "recording.flac"
        print("Preparing the complete recording as lossless FLAC for Pyannote...", file=sys.stderr)
        command = ["ffmpeg", "-hide_banner", "-loglevel", "error", "-nostdin", "-y",
                   "-i", str(audio), "-map", "0:a:0", "-vn", "-c:a", "flac",
                   "-compression_level", "5", str(flac)]
        try:
            encoded = subprocess.run(command, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE,
                                     timeout=client.remaining(), check=False)
        except subprocess.TimeoutExpired as error:
            raise ValueError("Pyannote audio encoding timed out") from error
        if encoded.returncode or not flac.is_file() or flac.stat().st_size == 0:
            raise ValueError("Could not encode the complete Pyannote recording as FLAC")
        if flac.stat().st_size > MAX_UPLOAD_BYTES:
            raise ValueError("Pyannote FLAC upload exceeds the 1 GiB limit")
        media_url = f"media://{uuid.uuid4()}.flac"
        media = client.request("POST", f"{API_URL}/media/input", payload={"url": media_url})
        signed_url = media.get("url")
        if not isinstance(signed_url, str):
            raise ValueError("Pyannote did not return an upload URL")
        print(f"Uploading the complete {duration / 60:.1f}-minute recording to Pyannote...", file=sys.stderr)
        client.request("PUT", signed_url, upload=flac, authenticated=False)
        payload = {"url": media_url, "model": model, "exclusive": True,
                   "turnLevelConfidence": True, "transcription": False}
        if speakers is not None:
            payload["numSpeakers"] = speakers
        created = client.request("POST", f"{API_URL}/diarize", payload=payload)
        warning = safe_detail({"message": created.get("warning")}, api_key)
        if warning:
            print(f"Warning: Pyannote: {warning}", file=sys.stderr)
        job_id = created.get("jobId")
        if not isinstance(job_id, str) or not job_id.strip() or any(ord(c) < 32 for c in job_id):
            raise ValueError("Pyannote did not return a valid job ID")
        print("Pyannote job submitted; waiting for whole-recording speaker detection...", file=sys.stderr)
        previous_status = None
        retries = 0
        while True:
            try:
                job = client.request("GET", f"{API_URL}/jobs/{quote(job_id, safe='')}")
            except RequestError as error:
                if not error.transient:
                    raise
                retries += 1
                print(f"Warning: {error}; retrying the same job (attempt {retries}).", file=sys.stderr)
                client.pause()
                continue
            if "jobId" in job and job["jobId"] != job_id:
                raise ValueError("Pyannote returned a different job ID while polling")
            status = job.get("status")
            if status == "succeeded":
                segments = validated_timeline(job)
                if any(segment["end"] > duration + MAX_NEAREST_GAP for segment in segments):
                    raise ValueError("Pyannote returned speaker times outside the uploaded recording")
                warning = safe_detail({"message": job["output"].get("warning")}, api_key)
                if warning:
                    print(f"Warning: Pyannote: {warning}", file=sys.stderr)
                retained = retained_result(job, job_id)
                completed = Path(temporary) / "completed.json"
                completed.write_text(json.dumps(retained, ensure_ascii=False) + "\n", encoding="utf-8")
                completed.replace(result_path)
                print("Pyannote whole-recording speaker detection succeeded.", file=sys.stderr)
                return retained
            if status in ("failed", "canceled"):
                detail = safe_detail(job, api_key) or safe_detail(job.get("output"), api_key)
                raise ValueError(f"Pyannote job {status}" + (f": {detail}" if detail else ""))
            if status not in ("created", "pending", "running"):
                raise ValueError("Pyannote job returned an unknown or missing status")
            if status != previous_status:
                print(f"Pyannote job status: {status}.", file=sys.stderr)
                previous_status = status
            client.pause()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)
    runner = subparsers.add_parser("run")
    runner.add_argument("audio")
    runner.add_argument("work")
    runner.add_argument("--model", default="precision-2")
    runner.add_argument("--speakers", type=int)
    runner.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT)
    apply = subparsers.add_parser("apply")
    apply.add_argument("transcript")
    apply.add_argument("result")
    apply.add_argument("--speakers", type=int)
    output = apply.add_mutually_exclusive_group()
    output.add_argument("--srt", action="store_true")
    output.add_argument("--json", action="store_true")
    args = parser.parse_args()
    try:
        if args.command == "run":
            run(args.audio, args.work, args.model, args.speakers, args.timeout)
            return
        response = align(read_json(args.transcript), read_json(args.result), args.speakers)
        if args.json:
            rendered = json.dumps(response, ensure_ascii=False) + "\n"
        else:
            rendered = convert(response) if args.srt else render_text(response)
    except (ValueError, OSError, wave.Error, EOFError) as error:
        print(f"Speaker detection failed: {error}", file=sys.stderr)
        sys.exit(1)
    sys.stdout.write(rendered)


if __name__ == "__main__":
    main()