#!/usr/bin/env python3
"""ffprobe wrapper that returns the facts an edit actually depends on.

Notable details other scripts kept getting wrong:
  * rotation. A phone clip is often stored 1920x1080 with a -90° display matrix.
    `width`/`height` from the stream are then the wrong way round, and a crop to
    9:16 silently destroys the frame. `display_width`/`display_height` here are
    post-rotation.
  * duration. Container duration and stream duration disagree on VFR phone
    footage; the honest answer is the largest reliable value, and both are kept.
  * frame rate. `r_frame_rate` is a bound, not a measurement, so `avg_frame_rate`
    is reported next to it and variable-rate footage is flagged.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
from fractions import Fraction
from pathlib import Path

from skill_config import emit, require_binary, run, utf8_stdout

AUDIO_EXTENSIONS = {".mp3", ".wav", ".m4a", ".aac", ".flac", ".ogg", ".opus", ".wma"}
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff", ".heic", ".avif"}
VIDEO_EXTENSIONS = {".mp4", ".mov", ".mkv", ".avi", ".webm", ".m4v", ".mpg", ".mpeg", ".wmv", ".flv", ".mts"}


def parse_rate(value: str | None) -> float | None:
    if not value or value in ("0/0", "N/A"):
        return None
    try:
        rate = float(Fraction(value))
    except (ValueError, ZeroDivisionError):
        return None
    return rate if rate > 0 else None


def rotation_of(stream: dict) -> int:
    """Effective clockwise display rotation in degrees."""
    for side_data in stream.get("side_data_list", []) or []:
        if "rotation" in side_data:
            try:
                return int(round(float(side_data["rotation"]))) % 360
            except (TypeError, ValueError):
                continue
    tag = (stream.get("tags") or {}).get("rotate")
    if tag:
        try:
            return int(round(float(tag))) % 360
        except ValueError:
            return 0
    return 0


def sha256_file(path: Path, chunk: int = 1 << 20) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while True:
            block = handle.read(chunk)
            if not block:
                break
            digest.update(block)
    return digest.hexdigest()


def probe(path: Path, with_hash: bool = False, count_frames: bool = False) -> dict:
    require_binary("ffprobe")
    path = Path(path)
    if not path.is_file():
        raise SystemExit(f"Not a file: {path}")
    command = [
        "ffprobe", "-v", "error", "-of", "json",
        "-show_format", "-show_streams", "-show_chapters",
    ]
    if count_frames:
        command += ["-count_frames"]
    command += [str(path)]
    raw = json.loads(run(command).stdout or "{}")
    streams = raw.get("streams", [])
    fmt = raw.get("format", {})
    video = next((s for s in streams if s.get("codec_type") == "video"
                  and (s.get("disposition", {}) or {}).get("attached_pic", 0) != 1), None)
    audio = next((s for s in streams if s.get("codec_type") == "audio"), None)

    container_duration = None
    try:
        container_duration = float(fmt.get("duration"))
    except (TypeError, ValueError):
        container_duration = None
    stream_durations = []
    for stream in streams:
        try:
            stream_durations.append(float(stream.get("duration")))
        except (TypeError, ValueError):
            continue
    duration = max([value for value in [container_duration, *stream_durations] if value], default=None)

    suffix = path.suffix.lower()
    if video is None and audio is not None:
        kind = "audio"
    elif suffix in IMAGE_EXTENSIONS or (video is not None and duration in (None, 0) and audio is None):
        kind = "image"
    elif video is not None:
        kind = "video"
    elif suffix in AUDIO_EXTENSIONS:
        kind = "audio"
    else:
        kind = "unknown"

    report: dict = {
        "path": str(path.resolve()),
        "name": path.name,
        "kind": kind,
        "bytes": path.stat().st_size,
        "container": fmt.get("format_name"),
        "duration": round(duration, 4) if duration else None,
        "container_duration": round(container_duration, 4) if container_duration else None,
        "bit_rate": int(fmt.get("bit_rate")) if str(fmt.get("bit_rate", "")).isdigit() else None,
        "has_video": video is not None,
        "has_audio": audio is not None,
        "streams": len(streams),
    }
    if with_hash:
        report["sha256"] = sha256_file(path)

    def stream_duration(stream: dict | None) -> float | None:
        if stream is None:
            return None
        try:
            return round(float(stream.get("duration")), 6)
        except (TypeError, ValueError):
            return None

    if video is not None:
        width = int(video.get("width") or 0)
        height = int(video.get("height") or 0)
        rotation = rotation_of(video)
        swapped = rotation in (90, 270)
        display_width, display_height = (height, width) if swapped else (width, height)
        r_rate = parse_rate(video.get("r_frame_rate"))
        avg_rate = parse_rate(video.get("avg_frame_rate"))
        sar = video.get("sample_aspect_ratio") or "1:1"
        report["video"] = {
            # Per-stream duration matters in a multi-step chain: trimming audio to
            # the *container* duration (which is the max of the two) makes the audio
            # a little longer at every pass, and the drift accumulates.
            "duration": stream_duration(video),
            "codec": video.get("codec_name"),
            "profile": video.get("profile"),
            "pix_fmt": video.get("pix_fmt"),
            "stored_width": width,
            "stored_height": height,
            "rotation": rotation,
            "display_width": display_width,
            "display_height": display_height,
            "display_aspect": round(display_width / display_height, 4) if display_height else None,
            "orientation": (
                "portrait" if display_height > display_width
                else "landscape" if display_width > display_height else "square"
            ),
            "r_frame_rate": round(r_rate, 4) if r_rate else None,
            "avg_frame_rate": round(avg_rate, 4) if avg_rate else None,
            "variable_frame_rate": bool(
                r_rate and avg_rate and abs(r_rate - avg_rate) / max(r_rate, avg_rate) > 0.02
            ),
            "sample_aspect_ratio": sar,
            "non_square_pixels": sar not in ("1:1", "0:1", "N/A"),
            "frames": int(video["nb_read_frames"]) if count_frames and str(
                video.get("nb_read_frames", "")).isdigit() else (
                int(video["nb_frames"]) if str(video.get("nb_frames", "")).isdigit() else None),
            "color_space": video.get("color_space"),
            "color_transfer": video.get("color_transfer"),
            "color_primaries": video.get("color_primaries"),
            "hdr": video.get("color_transfer") in ("smpte2084", "arib-std-b67"),
        }
    if audio is not None:
        report["audio"] = {
            "duration": stream_duration(audio),
            "codec": audio.get("codec_name"),
            "sample_rate": int(audio.get("sample_rate") or 0),
            "channels": int(audio.get("channels") or 0),
            "channel_layout": audio.get("channel_layout"),
            "bit_rate": int(audio.get("bit_rate")) if str(audio.get("bit_rate", "")).isdigit() else None,
        }
    return report


def loudness(path: Path) -> dict:
    """Measure integrated loudness, true peak and range with ebur128."""
    require_binary("ffmpeg")
    completed = run(
        ["ffmpeg", "-hide_banner", "-nostats", "-i", str(path),
         "-map", "0:a:0", "-filter:a", "ebur128=peak=true", "-f", "null", "-"],
        check=False,
    )
    text = completed.stderr or ""
    if "Summary:" not in text:
        return {"measured": False, "reason": "no audio stream or ebur128 unavailable"}
    tail = text[text.rindex("Summary:"):]
    numbers: dict[str, float] = {}
    for key, label in (
        ("i", "I:"), ("threshold", "Threshold:"), ("lra", "LRA:"),
        ("lra_low", "LRA low:"), ("lra_high", "LRA high:"),
        ("peak", "Peak:"),
    ):
        index = tail.find(label)
        if index < 0:
            continue
        chunk = tail[index + len(label):].split()
        for token in chunk[:2]:
            try:
                numbers[key] = float(token)
                break
            except ValueError:
                continue
    return {
        "measured": True,
        "integrated_lufs": numbers.get("i"),
        "true_peak_dbtp": numbers.get("peak"),
        "range_lu": numbers.get("lra"),
        "threshold_lufs": numbers.get("threshold"),
    }


def silence_map(path: Path, threshold_db: float = -34.0, minimum: float = 0.35) -> list[dict]:
    """Where the speech actually stops — used for pause-safe recuts."""
    require_binary("ffmpeg")
    completed = run(
        ["ffmpeg", "-hide_banner", "-nostats", "-i", str(path), "-map", "0:a:0",
         "-af", f"silencedetect=noise={threshold_db}dB:d={minimum}", "-f", "null", "-"],
        check=False,
    )
    spans: list[dict] = []
    start: float | None = None
    for line in (completed.stderr or "").splitlines():
        if "silence_start:" in line:
            try:
                start = float(line.split("silence_start:")[1].split()[0])
            except (IndexError, ValueError):
                start = None
        elif "silence_end:" in line and start is not None:
            try:
                end = float(line.split("silence_end:")[1].split()[0])
            except (IndexError, ValueError):
                continue
            spans.append({"start": round(start, 3), "end": round(end, 3),
                          "duration": round(end - start, 3)})
            start = None
    return spans


def fit_transform(source: dict, target_width: int, target_height: int) -> dict:
    """How to place `source` on the target canvas without ever stretching it.

    Returns the scale/crop numbers plus how much of the source is lost, so the
    caller can decide between a crop and a blurred-pad before rendering.
    """
    video = source.get("video") or {}
    width = video.get("display_width") or target_width
    height = video.get("display_height") or target_height
    source_aspect = width / height
    target_aspect = target_width / target_height
    if source_aspect > target_aspect:
        scale = target_height / height
        cropped_fraction = 1 - target_aspect / source_aspect
        axis = "horizontal"
    else:
        scale = target_width / width
        cropped_fraction = 1 - source_aspect / target_aspect
        axis = "vertical"
    upscale = scale > 1.0
    return {
        "scale": round(scale, 5),
        "cover_filter": (
            f"scale={target_width}:{target_height}:force_original_aspect_ratio=increase:flags=lanczos,"
            f"crop={target_width}:{target_height}"
        ),
        "pad_filter": (
            f"scale={target_width}:{target_height}:force_original_aspect_ratio=increase:flags=lanczos,"
            f"crop={target_width}:{target_height},gblur=sigma=28[bg];"
            f"scale={target_width}:-2:flags=lanczos[fg];[bg][fg]overlay=(W-w)/2:(H-h)/2"
        ),
        "cropped_axis": axis,
        "cropped_fraction": round(cropped_fraction, 4),
        "upscaled": upscale,
        "upscale_factor": round(scale, 3) if upscale else 1.0,
        "recommendation": (
            "pad" if cropped_fraction > 0.42 else
            "cover_with_review" if cropped_fraction > 0.28 else "cover"
        ),
        "warning": (
            "Кроп срезает больше 42% кадра: возьми другой исходник или используй "
            "размытый паддинг." if cropped_fraction > 0.42 else
            f"Апскейл x{scale:.2f} — проверь резкость на телефоне."
            if scale > 1.35 else None
        ),
    }


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("media", nargs="+")
    parser.add_argument("--hash", action="store_true", help="Add SHA-256 (slow on large files)")
    parser.add_argument("--count-frames", action="store_true", help="Decode to count frames exactly")
    parser.add_argument("--loudness", action="store_true", help="Measure EBU R128 loudness")
    parser.add_argument("--silence", action="store_true", help="Report silent spans")
    parser.add_argument("--fit", help="Target canvas WxH, e.g. 1080x1920")
    parser.add_argument("--report", help="Write JSON here")
    args = parser.parse_args()

    results = []
    for item in args.media:
        report = probe(Path(item), with_hash=args.hash, count_frames=args.count_frames)
        if args.loudness and report.get("has_audio"):
            report["loudness"] = loudness(Path(item))
        if args.silence and report.get("has_audio"):
            report["silence"] = silence_map(Path(item))
        if args.fit and report.get("has_video"):
            width, _, height = args.fit.partition("x")
            report["fit"] = fit_transform(report, int(width), int(height))
        results.append(report)

    payload = {"count": len(results), "media": results}
    if args.report:
        destination = Path(args.report)
        destination.parent.mkdir(parents=True, exist_ok=True)
        destination.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    emit(payload)


if __name__ == "__main__":
    main()