import argparse
import json
import math
from pathlib import Path
from typing import Any


def is_nonfinite_number(value: Any) -> bool:
    """Prüft auf NaN sowie positive oder negative Unendlichkeit."""
    return isinstance(value, float) and not math.isfinite(value)


def sanitize_nonfinite_values(value: Any) -> Any:
    """
    Ersetzt NaN und Infinity rekursiv durch None.

    None wird in JSON als null geschrieben.
    """
    if is_nonfinite_number(value):
        return None

    if isinstance(value, dict):
        return {
            key: sanitize_nonfinite_values(item)
            for key, item in value.items()
        }

    if isinstance(value, list):
        return [
            sanitize_nonfinite_values(item)
            for item in value
        ]

    return value


def is_invalid_empty_segment(segment: dict[str, Any]) -> bool:
    """Erkennt technisch leere Whisper-Artefakte."""

    text = str(segment.get("text", "")).strip()
    start = segment.get("start")
    end = segment.get("end")
    avg_logprob = segment.get("avg_logprob")

    logprob_is_nan = (
        isinstance(avg_logprob, float)
        and math.isnan(avg_logprob)
    )

    return (
        not text
        and start == end
        and logprob_is_nan
    )


def clean_whisper_json(
    input_path: Path,
    output_path: Path,
) -> None:
    with input_path.open("r", encoding="utf-8") as file:
        data = json.load(file)

    original_segments = data.get("segments", [])

    cleaned_segments = [
        segment
        for segment in original_segments
        if not is_invalid_empty_segment(segment)
    ]

    for new_id, segment in enumerate(cleaned_segments):
        segment["id"] = new_id

    data["segments"] = cleaned_segments

    data["text"] = " ".join(
        str(segment.get("text", "")).strip()
        for segment in cleaned_segments
        if str(segment.get("text", "")).strip()
    )

    # Alle noch vorhandenen NaN-/Infinity-Werte durch null ersetzen.
    data = sanitize_nonfinite_values(data)

    # Erst vollständig in einen String serialisieren.
    # Dadurch bleibt keine unvollständige Ausgabedatei zurück,
    # falls doch noch ein Fehler auftritt.
    json_content = json.dumps(
        data,
        ensure_ascii=False,
        indent=2,
        allow_nan=False,
    )

    output_path.write_text(
        json_content + "\n",
        encoding="utf-8",
    )

    removed = len(original_segments) - len(cleaned_segments)

    print(f"Eingabedatei:      {input_path}")
    print(f"Ausgabedatei:      {output_path}")
    print(f"Segmente vorher:   {len(original_segments)}")
    print(f"Segmente nachher:  {len(cleaned_segments)}")
    print(f"Segmente entfernt: {removed}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Entfernt technisch leere Artefakte aus "
            "einer MLX-Whisper-JSON-Datei."
        )
    )
    parser.add_argument(
        "input",
        type=Path,
        help="Whisper-JSON-Datei",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        help="Ausgabedatei",
    )

    args = parser.parse_args()

    output_path = args.output or args.input.with_name(
        f"{args.input.stem}_cleaned.json"
    )

    clean_whisper_json(args.input, output_path)


if __name__ == "__main__":
    main()