from __future__ import annotations

import json
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from scipy import signal

MARKER_FILENAME = "marqueur_pub.wav"
DECODE_SAMPLE_RATE = 2000
CORRELATION_THRESHOLD = 0.80
MIN_PEAK_DISTANCE_SECONDS = 30.0
MAX_AD_DURATION_SECONDS = 300.0
DEFAULT_BITRATE = "128k"
EPSILON = 1e-9


class PodcastCleanerError(RuntimeError):
    """Raised when ad detection or audio cutting cannot be completed."""


@dataclass(frozen=True)
class MarkerHit:
    start_seconds: float
    end_seconds: float
    score: float


@dataclass(frozen=True)
class AdSegment:
    start_seconds: float
    end_seconds: float

    @property
    def duration_seconds(self) -> float:
        return self.end_seconds - self.start_seconds


@dataclass(frozen=True)
class AudioMetadata:
    duration_seconds: float
    bitrate: str


@dataclass(frozen=True)
class CleanResult:
    input_path: Path
    output_audio_path: Path
    report_path: Path
    marker_hits: tuple[MarkerHit, ...]
    ad_segments: tuple[AdSegment, ...]
    input_metadata: AudioMetadata
    output_metadata: AudioMetadata


def decode_audio_for_detection(audio_path: Path, sample_rate: int = DECODE_SAMPLE_RATE) -> np.ndarray:
    command = [
        "ffmpeg",
        "-v",
        "error",
        "-i",
        str(audio_path),
        "-ac",
        "1",
        "-ar",
        str(sample_rate),
        "-f",
        "s16le",
        "-",
    ]

    try:
        completed = subprocess.run(command, capture_output=True, check=True)
    except FileNotFoundError as exc:
        raise PodcastCleanerError("ffmpeg is not installed or not available in PATH.") from exc
    except subprocess.CalledProcessError as exc:
        stderr = exc.stderr.decode("utf-8", errors="replace").strip()
        raise PodcastCleanerError(f"Unable to decode audio for detection: {stderr or exc}") from exc

    samples = np.frombuffer(completed.stdout, dtype=np.int16).astype(np.float32)
    if samples.size == 0:
        raise PodcastCleanerError(f"No audio samples were decoded from {audio_path}.")
    return samples / 32768.0


def detect_marker_hits(
    marker_samples: np.ndarray,
    podcast_samples: np.ndarray,
    sample_rate: int = DECODE_SAMPLE_RATE,
    threshold: float = CORRELATION_THRESHOLD,
    min_peak_distance_seconds: float = MIN_PEAK_DISTANCE_SECONDS,
) -> tuple[list[MarkerHit], float]:
    if podcast_samples.size < marker_samples.size:
        raise PodcastCleanerError("The podcast audio is shorter than the marker audio.")

    marker_signal = marker_samples.astype(np.float32, copy=True)
    podcast_signal = podcast_samples.astype(np.float32, copy=True)
    marker_signal -= marker_signal.mean()
    podcast_signal -= podcast_signal.mean()

    marker_energy = float(np.sum(marker_signal ** 2))
    if marker_energy <= EPSILON:
        raise PodcastCleanerError("The marker audio is silent or invalid.")

    correlation = signal.fftconvolve(podcast_signal, marker_signal[::-1], mode="valid")
    window_energy = signal.fftconvolve(
        podcast_signal ** 2,
        np.ones(marker_signal.size, dtype=np.float32),
        mode="valid",
    )
    denominator = np.sqrt(np.maximum(window_energy * marker_energy, EPSILON))
    scores = correlation / denominator

    minimum_distance = max(1, int(min_peak_distance_seconds * sample_rate))
    peaks, properties = signal.find_peaks(scores, height=threshold, distance=minimum_distance)

    marker_duration_seconds = marker_signal.size / sample_rate
    hits = [
        MarkerHit(
            start_seconds=index / sample_rate,
            end_seconds=(index / sample_rate) + marker_duration_seconds,
            score=float(properties["peak_heights"][offset]),
        )
        for offset, index in enumerate(peaks)
    ]
    return hits, marker_duration_seconds


def validate_and_pair_markers(
    marker_hits: list[MarkerHit],
    marker_duration_seconds: float,
    max_ad_duration_seconds: float = MAX_AD_DURATION_SECONDS,
) -> list[AdSegment]:
    if not marker_hits:
        raise PodcastCleanerError("No marker occurrences were detected in the podcast.")

    if len(marker_hits) % 2 != 0:
        raise PodcastCleanerError(f"Odd marker count detected: {len(marker_hits)}.")

    ad_segments: list[AdSegment] = []
    for index in range(0, len(marker_hits), 2):
        first_hit = marker_hits[index]
        second_hit = marker_hits[index + 1]
        ad_segment = AdSegment(
            start_seconds=first_hit.start_seconds,
            end_seconds=second_hit.start_seconds + marker_duration_seconds,
        )
        if ad_segment.duration_seconds > max_ad_duration_seconds:
            raise PodcastCleanerError(
                "Detected ad segment exceeds 300 seconds: "
                f"{ad_segment.start_seconds:.3f} -> {ad_segment.end_seconds:.3f}"
            )
        ad_segments.append(ad_segment)

    return ad_segments


def compute_keep_segments(total_duration_seconds: float, ad_segments: list[AdSegment]) -> list[AdSegment]:
    keep_segments: list[AdSegment] = []
    cursor = 0.0

    for ad_segment in ad_segments:
        if ad_segment.start_seconds > cursor + EPSILON:
            keep_segments.append(AdSegment(start_seconds=cursor, end_seconds=ad_segment.start_seconds))
        cursor = max(cursor, ad_segment.end_seconds)

    if total_duration_seconds > cursor + EPSILON:
        keep_segments.append(AdSegment(start_seconds=cursor, end_seconds=total_duration_seconds))

    return [segment for segment in keep_segments if segment.duration_seconds > 0.05]


def probe_audio_metadata(audio_path: Path) -> AudioMetadata:
    command = [
        "ffprobe",
        "-v",
        "error",
        "-show_entries",
        "format=duration,bit_rate",
        "-of",
        "json",
        str(audio_path),
    ]

    try:
        completed = subprocess.run(command, capture_output=True, text=True, check=True)
    except FileNotFoundError as exc:
        raise PodcastCleanerError("ffprobe is not installed or not available in PATH.") from exc
    except subprocess.CalledProcessError as exc:
        stderr = exc.stderr.strip()
        raise PodcastCleanerError(f"Unable to probe audio metadata: {stderr or exc}") from exc

    try:
        payload = json.loads(completed.stdout)
    except json.JSONDecodeError as exc:
        raise PodcastCleanerError(f"Unable to parse ffprobe output: {exc}") from exc

    format_info = payload.get("format", {})
    duration_value = format_info.get("duration")
    if duration_value is None:
        raise PodcastCleanerError("ffprobe did not return the podcast duration.")

    try:
        duration_seconds = float(duration_value)
    except ValueError as exc:
        raise PodcastCleanerError(f"Invalid duration reported by ffprobe: {duration_value!r}") from exc

    bitrate_value = format_info.get("bit_rate")
    bitrate = DEFAULT_BITRATE
    if bitrate_value:
        try:
            bitrate = f"{int(bitrate_value) // 1000}k"
        except ValueError:
            bitrate = DEFAULT_BITRATE

    return AudioMetadata(duration_seconds=duration_seconds, bitrate=bitrate)


def build_audacity_report_lines(
    marker_hits: list[MarkerHit],
    ad_segments: list[AdSegment],
    errors: list[str] | None = None,
) -> list[str]:
    lines: list[str] = []

    for index, marker_hit in enumerate(marker_hits, start=1):
        lines.append(
            f"{marker_hit.start_seconds:.6f}\t{marker_hit.end_seconds:.6f}\t"
            f"marker_{index:02d} score={marker_hit.score:.3f}"
        )

    for index, ad_segment in enumerate(ad_segments, start=1):
        lines.append(f"{ad_segment.start_seconds:.6f}\t{ad_segment.end_seconds:.6f}\tpub_{index:02d}")

    for error in errors or []:
        lines.append(f"0.000000\t0.000000\tERROR {error}")

    return lines


def write_audacity_report(
    report_path: Path,
    marker_hits: list[MarkerHit],
    ad_segments: list[AdSegment],
    errors: list[str] | None = None,
) -> None:
    report_path.parent.mkdir(parents=True, exist_ok=True)
    lines = build_audacity_report_lines(marker_hits, ad_segments, errors)
    report_path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def run_ffmpeg_cut(
    input_path: Path,
    output_path: Path,
    keep_segments: list[AdSegment],
    bitrate: str,
) -> None:
    if not keep_segments:
        raise PodcastCleanerError("No audio remains after removing the detected ad segments.")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    filter_parts: list[str] = []
    concat_inputs: list[str] = []
    for index, segment in enumerate(keep_segments):
        filter_parts.append(
            f"[0:a]atrim=start={segment.start_seconds:.6f}:end={segment.end_seconds:.6f},"
            f"asetpts=PTS-STARTPTS[seg{index}]"
        )
        concat_inputs.append(f"[seg{index}]")

    filter_parts.append(f"{''.join(concat_inputs)}concat=n={len(keep_segments)}:v=0:a=1[outa]")
    filter_complex = ";".join(filter_parts)
    temp_output = output_path.with_name(f"{output_path.stem}.part{output_path.suffix}")

    command = [
        "ffmpeg",
        "-y",
        "-v",
        "error",
        "-i",
        str(input_path),
        "-filter_complex",
        filter_complex,
        "-map",
        "[outa]",
        "-c:a",
        "libmp3lame",
        "-b:a",
        bitrate,
        str(temp_output),
    ]

    try:
        subprocess.run(command, capture_output=True, text=True, check=True)
        temp_output.replace(output_path)
    except FileNotFoundError as exc:
        raise PodcastCleanerError("ffmpeg is not installed or not available in PATH.") from exc
    except subprocess.CalledProcessError as exc:
        stderr = exc.stderr.strip()
        raise PodcastCleanerError(f"Unable to generate cleaned podcast: {stderr or exc}") from exc
    finally:
        if temp_output.exists():
            temp_output.unlink()


def build_output_paths(input_path: Path) -> tuple[Path, Path]:
    output_audio_path = input_path.with_name(f"{input_path.stem}.sans_pub.mp3")
    report_path = input_path.with_name(f"{input_path.stem}.pubs_audacity.txt")
    return output_audio_path, report_path


def clean_podcast(
    input_path: Path,
    output_audio_path: Path | None = None,
    report_path: Path | None = None,
    marker_path: Path | None = None,
) -> CleanResult:
    input_path = Path(input_path).expanduser().resolve()
    output_audio_path = Path(output_audio_path).resolve() if output_audio_path else build_output_paths(input_path)[0]
    report_path = Path(report_path).resolve() if report_path else build_output_paths(input_path)[1]
    marker_path = Path(marker_path).resolve() if marker_path else Path(__file__).resolve().with_name(MARKER_FILENAME)

    marker_hits: list[MarkerHit] = []
    ad_segments: list[AdSegment] = []
    error_messages: list[str] = []

    try:
        if not input_path.is_file():
            raise PodcastCleanerError(f"Podcast file not found: {input_path}")
        if not marker_path.is_file():
            raise PodcastCleanerError(f"Marker file not found: {marker_path}")

        input_metadata = probe_audio_metadata(input_path)
        marker_samples = decode_audio_for_detection(marker_path)
        podcast_samples = decode_audio_for_detection(input_path)
        marker_hits, marker_duration_seconds = detect_marker_hits(marker_samples, podcast_samples)
        ad_segments = validate_and_pair_markers(marker_hits, marker_duration_seconds)
        keep_segments = compute_keep_segments(input_metadata.duration_seconds, ad_segments)

        write_audacity_report(report_path, marker_hits, ad_segments)
        run_ffmpeg_cut(input_path, output_audio_path, keep_segments, input_metadata.bitrate)
        output_metadata = probe_audio_metadata(output_audio_path)
        return CleanResult(
            input_path=input_path,
            output_audio_path=output_audio_path,
            report_path=report_path,
            marker_hits=tuple(marker_hits),
            ad_segments=tuple(ad_segments),
            input_metadata=input_metadata,
            output_metadata=output_metadata,
        )
    except PodcastCleanerError as exc:
        error_messages.append(str(exc))
        write_audacity_report(report_path, marker_hits, ad_segments, error_messages)
        raise


def main(argv: list[str] | None = None) -> int:
    args = list(sys.argv[1:] if argv is None else argv)
    if len(args) != 1:
        print("Usage: python clean_podcast_ads.py <podcast.mp3>", file=sys.stderr)
        return 1

    input_path = Path(args[0]).expanduser().resolve()
    output_audio_path, report_path = build_output_paths(input_path)

    try:
        result = clean_podcast(input_path, output_audio_path=output_audio_path, report_path=report_path)
        print(f"Markers detected: {len(result.marker_hits)}")
        print(f"Ad segments detected: {len(result.ad_segments)}")
        for index, ad_segment in enumerate(result.ad_segments, start=1):
            print(
                f"pub_{index:02d}: {ad_segment.start_seconds:.3f}s -> "
                f"{ad_segment.end_seconds:.3f}s ({ad_segment.duration_seconds:.3f}s)"
            )
        print(f"Audacity report: {result.report_path}")
        print(f"Cleaned MP3: {result.output_audio_path}")
        return 0
    except PodcastCleanerError as exc:
        print(f"Audacity report: {report_path}")
        print(f"Error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())