#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import json
import math
import os
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np

MARKER_TAG = "automation-agent:slip-clips-to-match-reference"
LOW_DIM = 480
HIGH_DIM = 1280
SHORTLIST_REGIONS = 12
SHORTLIST_RADIUS_FRAMES = 2
PREMIERE_TICKS_PER_SECOND = 254_016_000_000
PREMIERE_TICK_SECONDS = 1.0 / PREMIERE_TICKS_PER_SECOND


def format_duration(seconds: float) -> str:
    milliseconds = max(0, round(float(seconds) * 1000))
    hours, remainder = divmod(milliseconds, 3_600_000)
    minutes, remainder = divmod(remainder, 60_000)
    whole_seconds, millis = divmod(remainder, 1000)
    return f"{hours:02d}:{minutes:02d}:{whole_seconds:02d}.{millis:03d}"


def update_processing_time_report(
    path: Path,
    clip_no: int | None,
    results: list[dict],
    elapsed_seconds: float,
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    existing_lines = path.read_text(encoding="utf-8").splitlines() if path.is_file() else []
    clip_lines = [line for line in existing_lines if line.startswith("Clip ")]
    matched = sum(1 for result in results if result.get("status") == "MATCHED")
    no_match = len(results) - matched
    label = str(clip_no) if clip_no is not None else "all"
    clip_lines.append(
        f"Clip {label} | analyzed {len(results)} | matched {matched} | no match {no_match} | "
        f"{elapsed_seconds:.3f} seconds | {format_duration(elapsed_seconds)}"
    )
    cumulative_seconds = 0.0
    for line in clip_lines:
        fields = line.split(" | ")
        if len(fields) >= 2 and fields[-2].endswith(" seconds"):
            cumulative_seconds += float(fields[-2].removesuffix(" seconds"))
    content = [
        "Automation Agent - Slip Clips To Match Reference",
        "",
        *clip_lines,
        "",
        f"Cumulative processing time: {format_duration(cumulative_seconds)}",
        f"Cumulative processing time seconds: {cumulative_seconds:.3f}",
        "",
    ]
    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
    try:
        temporary.write_text("\n".join(content), encoding="utf-8")
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def warn_processing_time_report(path: Path, exc: Exception) -> None:
    message = str(exc).replace("|", "/").replace("\r", " ").replace("\n", " ")
    safe_path = str(path).replace("|", "/").replace("\r", " ").replace("\n", " ")
    print(
        f"WARNING|stage=processing-time-report|message={message}|path={safe_path}",
        file=sys.stderr,
        flush=True,
    )


@dataclass
class ClipRow:
    n: int
    ref: str
    ref_clip_ref: str
    start: float
    end: float
    in_point: float
    out_point: float
    src: str
    ref_start: float
    ref_end: float
    ref_in: float
    ref_src: str
    source_fps: float
    media_duration: float
    selected: bool

    @property
    def duration(self) -> float:
        return max(0.0, self.end - self.start)

    @property
    def overlap_start(self) -> float:
        return max(self.start, self.ref_start)

    @property
    def overlap_end(self) -> float:
        return min(self.end, self.ref_end)


def run(cmd: list[str], timeout_seconds: float) -> subprocess.CompletedProcess:
    result = subprocess.run(
        cmd,
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        timeout=timeout_seconds,
    )
    if result.returncode:
        raise RuntimeError(
            f"External command failed ({result.returncode}): {cmd}\n"
            + result.stderr.decode("utf-8", errors="replace")
        )
    return result


def export_frame(ffmpeg: str, media: str, when: float, output: Path, max_dim: int, timeout_seconds: float) -> np.ndarray:
    output.parent.mkdir(parents=True, exist_ok=True)
    output.unlink(missing_ok=True)
    run([
        ffmpeg, "-y", "-ss", f"{max(0.0, when):.9f}", "-i", media,
        "-frames:v", "1", "-vf",
        f"scale='if(gt(iw,ih),min({max_dim},iw),-2)':'if(gt(iw,ih),-2,min({max_dim},ih))'",
        "-update", "1", str(output),
    ], timeout_seconds)
    image = cv2.imread(str(output), cv2.IMREAD_GRAYSCALE)
    if image is None:
        raise RuntimeError(f"Could not read extracted frame: {output}")
    return image


def export_candidate_series(
    ffmpeg: str,
    media: str,
    first_time: float,
    count: int,
    output_fps: float,
    output_dir: Path,
    prefix: str,
    max_dim: int = LOW_DIM,
    timeout_seconds: float = 600.0,
) -> list[np.ndarray]:
    """Decode an evenly-spaced candidate pass with one FFmpeg process."""
    output_dir.mkdir(parents=True, exist_ok=True)
    for previous in output_dir.glob(f"{prefix}_*.jpg"):
        previous.unlink(missing_ok=True)
    pattern = output_dir / f"{prefix}_%06d.jpg"
    duration = max(1.0 / output_fps, (count + 1) / output_fps)
    run([
        ffmpeg, "-y", "-ss", f"{max(0.0, first_time):.9f}", "-i", media,
        "-t", f"{duration:.9f}", "-vf",
        f"fps={output_fps:.12f},scale='if(gt(iw,ih),min({max_dim},iw),-2)':'if(gt(iw,ih),-2,min({max_dim},ih))'",
        "-frames:v", str(count), "-q:v", "4", str(pattern),
    ], timeout_seconds)
    images = []
    for image_path in sorted(output_dir.glob(f"{prefix}_*.jpg")):
        image = cv2.imread(str(image_path), cv2.IMREAD_GRAYSCALE)
        if image is not None:
            images.append(image)
    return images


def contiguous_groups(values: list[int]) -> list[list[int]]:
    groups: list[list[int]] = []
    for value in sorted(values):
        if not groups or value != groups[-1][-1] + 1:
            groups.append([value])
        else:
            groups[-1].append(value)
    return groups


def normalize(image: np.ndarray, size: int = 96) -> np.ndarray:
    h, w = image.shape[:2]
    scale = size / max(h, w)
    resized = cv2.resize(image, (max(8, round(w * scale)), max(8, round(h * scale))), interpolation=cv2.INTER_AREA)
    canvas = np.zeros((size, size), dtype=np.uint8)
    y = (size - resized.shape[0]) // 2
    x = (size - resized.shape[1]) // 2
    canvas[y:y + resized.shape[0], x:x + resized.shape[1]] = resized
    canvas = cv2.equalizeHist(canvas)
    return cv2.Sobel(canvas, cv2.CV_32F, 1, 1, ksize=3)


def cheap_score(a: np.ndarray, b: np.ndarray) -> float:
    aa = normalize(a).reshape(-1)
    bb = normalize(b).reshape(-1)
    aa -= aa.mean()
    bb -= bb.mean()
    denom = float(np.linalg.norm(aa) * np.linalg.norm(bb))
    return 0.0 if denom < 1e-6 else float(np.dot(aa, bb) / denom)


def feature_score(a: np.ndarray, b: np.ndarray) -> tuple[float, int, float]:
    detector = cv2.SIFT_create(nfeatures=5000) if hasattr(cv2, "SIFT_create") else cv2.AKAZE_create()
    ka, da = detector.detectAndCompute(a, None)
    kb, db = detector.detectAndCompute(b, None)
    if da is None or db is None or len(ka) < 8 or len(kb) < 8:
        return 0.0, 0, 0.0
    norm = cv2.NORM_L2 if da.dtype == np.float32 else cv2.NORM_HAMMING
    matches = cv2.BFMatcher(norm).knnMatch(da, db, k=2)
    good = [m[0] for m in matches if len(m) == 2 and m[0].distance < 0.74 * m[1].distance]
    if len(good) < 8:
        return 0.0, 0, 0.0
    pa = np.float32([ka[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
    pb = np.float32([kb[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
    _, mask = cv2.estimateAffinePartial2D(pa, pb, method=cv2.RANSAC, ransacReprojThreshold=5.0)
    if mask is None:
        return 0.0, 0, 0.0
    inliers = int(mask.sum())
    ratio = inliers / max(1, len(good))
    return min(1.0, (inliers / 40.0) * ratio), inliers, ratio


def distinct_shortlist(scores: list[tuple[int, float]], count: int) -> list[int]:
    chosen: list[int] = []
    for frame, _ in sorted(scores, key=lambda pair: pair[1], reverse=True):
        if all(abs(frame - prior) > SHORTLIST_RADIUS_FRAMES for prior in chosen):
            chosen.append(frame)
        if len(chosen) >= count:
            break
    return chosen


def anchors(row: ClipRow) -> list[float]:
    start, end = row.overlap_start, row.overlap_end
    duration = end - start
    if duration <= 0:
        return []
    if duration < 0.2:
        return [start + duration * 0.5]
    if duration < 0.75:
        return [start + duration * 0.25, start + duration * 0.75]
    return [start + duration * p for p in (0.2, 0.5, 0.8)]


def candidate_bounds(row: ClipRow, mode: str, custom_seconds: float) -> tuple[int, int]:
    fps = max(1.0, row.source_fps)
    backward = max(0.0, row.in_point)
    forward = max(0.0, row.media_duration - row.out_point)
    requested = {"3_SECONDS": 3.0, "10_SECONDS": 10.0, "CUSTOM": max(0.0, custom_seconds)}.get(mode)
    if requested is not None:
        backward = min(backward, requested)
        forward = min(forward, requested)
    return -math.floor(backward * fps + 1e-6), math.floor(forward * fps + 1e-6)


def candidate_frames(lo: int, hi: int, fps: float, entire: bool) -> list[int]:
    total = hi - lo + 1
    if not entire or total <= 1200:
        return list(range(lo, hi + 1))
    coarse_step = max(1, round(fps / 2.0))
    result = list(range(lo, hi + 1, coarse_step))
    if result[-1] != hi:
        result.append(hi)
    return result


def quality_label(best: float, runner_up: float, support: int, total: int) -> tuple[str, int]:
    margin = best - runner_up
    if best >= 0.62 and margin >= 0.10 and support == total and total >= 2:
        return "High", 2
    if best >= 0.38 and margin >= 0.035 and support >= max(1, total - 1):
        return "Medium", 3
    return "Low", 1


def same_media_path(left: str, right: str) -> bool:
    return os.path.normcase(os.path.abspath(left)) == os.path.normcase(os.path.abspath(right))


def phase_guarded_visual_range(
    in_point: float,
    out_point: float,
    offset_frames: int,
    source_fps: float,
) -> tuple[float, float, float]:
    """Return an idempotent visual-match range just inside its source frame.

    Premiere track-item starts can sit one tick after a sequence frame boundary.
    An in-point exactly on the corresponding source frame boundary can then show
    the preceding frame in Premiere even though FFmpeg resolves the expected
    source frame. Re-align to the matched source frame and move one Premiere
    tick inside it. Repeated runs retain the same guarded in-point instead of
    accumulating additional nudges.
    """
    source_frame = round(in_point * source_fps) + offset_frames
    guarded_in = source_frame / source_fps + PREMIERE_TICK_SECONDS
    applied_offset = guarded_in - in_point
    if abs(applied_offset) < PREMIERE_TICK_SECONDS / 2.0:
        applied_offset = 0.0
        guarded_in = in_point
    return guarded_in, out_point + applied_offset, applied_offset


def clean(value: object) -> str:
    return str(value).replace("\t", " ").replace("\r", " ").replace("\n", " | ")


def parse_manifest(path: Path) -> list[ClipRow]:
    rows: list[ClipRow] = []
    with path.open("r", encoding="utf-8", newline="") as handle:
        for raw in csv.DictReader(handle, delimiter="\t"):
            rows.append(ClipRow(
                n=int(raw["clipNo"]), ref=raw["clipRef"], ref_clip_ref=raw["refClipRef"],
                start=float(raw["start"]), end=float(raw["end"]), in_point=float(raw["in"]),
                out_point=float(raw["out"]), src=raw["src"], ref_start=float(raw["refStart"]),
                ref_end=float(raw["refEnd"]), ref_in=float(raw["refIn"]), ref_src=raw["refSrc"],
                source_fps=float(raw.get("sourceFps") or 30.0),
                media_duration=float(raw.get("mediaDuration") or raw["out"]),
                selected=str(raw.get("selected", "true")).lower() == "true",
            ))
    return rows


def analyze(row: ClipRow, args: argparse.Namespace, frames_dir: Path) -> dict:
    times = anchors(row)
    if not times:
        return {"status": "NO_MATCH", "reason": "no_overlap"}
    lo, hi = candidate_bounds(row, args.search_mode, args.custom_seconds)
    if lo == 0 and hi == 0:
        return {"status": "NO_MATCH", "reason": "no_source_handles", "lo": lo, "hi": hi}

    if same_media_path(row.src, row.ref_src):
        new_in = row.ref_in + (row.start - row.ref_start)
        applied_offset_seconds = new_in - row.in_point
        offset_frames = round(applied_offset_seconds * row.source_fps)
        new_out = row.out_point + applied_offset_seconds
        if offset_frames < lo or offset_frames > hi:
            return {
                "status": "NO_MATCH",
                "reason": "same_source_match_outside_search_range",
                "lo": lo,
                "hi": hi,
            }
        if new_in < -1e-7 or new_out > row.media_duration + 1e-7:
            return {
                "status": "NO_MATCH",
                "reason": "same_source_match_exceeds_handles",
                "lo": lo,
                "hi": hi,
            }
        return {
            "status": "MATCHED",
            "matchMethod": "SAME_SOURCE_TIMELINE_MAPPING",
            "offsetFrames": offset_frames,
            "offsetSeconds": applied_offset_seconds,
            "appliedOffsetSeconds": applied_offset_seconds,
            "newIn": new_in,
            "newOut": new_out,
            "score": 1.0,
            "runnerUp": 0.0,
            "support": len(times),
            "samples": len(times),
            "quality": "High",
            "color": 2,
            "lo": lo,
            "hi": hi,
            "coarseCandidates": 0,
            "verifiedCandidates": 1,
        }

    entire = args.search_mode == "ENTIRE_SOURCE"
    coarse = candidate_frames(lo, hi, row.source_fps, entire)
    reference_low = []
    for i, timeline_time in enumerate(times):
        ref_time = row.ref_in + (timeline_time - row.ref_start)
        reference_low.append(export_frame(args.ffmpeg, row.ref_src, ref_time, frames_dir / f"c{row.n}_ref_{i}_low.png", LOW_DIM, args.ffmpeg_timeout_seconds))

    scores: list[tuple[int, float]] = []
    probe_anchor = times[len(times) // 2]
    probe_ref = reference_low[len(times) // 2]
    coarse_step_frames = max(1, coarse[1] - coarse[0]) if len(coarse) > 1 else 1
    first_source_time = row.in_point + (probe_anchor - row.start) + coarse[0] / row.source_fps
    decoded = export_candidate_series(
        args.ffmpeg,
        row.src,
        first_source_time,
        len(coarse),
        row.source_fps / coarse_step_frames,
        frames_dir,
        f"c{row.n}_coarse",
        timeout_seconds=args.ffmpeg_timeout_seconds,
    )
    for offset_frame, source in zip(coarse, decoded):
        scores.append((offset_frame, cheap_score(source, probe_ref)))
    if not scores:
        return {"status": "NO_MATCH", "reason": "no_decodable_candidates", "lo": lo, "hi": hi}

    regions = distinct_shortlist(scores, SHORTLIST_REGIONS)
    fine: set[int] = set()
    fine_radius = max(SHORTLIST_RADIUS_FRAMES, round(row.source_fps)) if entire and len(coarse) < hi - lo + 1 else SHORTLIST_RADIUS_FRAMES
    for center in regions:
        fine.update(range(max(lo, center - fine_radius), min(hi, center + fine_radius) + 1))

    reference_high = []
    for i, timeline_time in enumerate(times):
        ref_time = row.ref_in + (timeline_time - row.ref_start)
        reference_high.append(export_frame(args.ffmpeg, row.ref_src, ref_time, frames_dir / f"c{row.n}_ref_{i}_high.png", HIGH_DIM, args.ffmpeg_timeout_seconds))

    verification_frames: dict[tuple[int, int], np.ndarray] = {}
    for sample_index, timeline_time in enumerate(times):
        for group_index, group in enumerate(contiguous_groups(sorted(fine))):
            first_time = row.in_point + (timeline_time - row.start) + group[0] / row.source_fps
            decoded_group = export_candidate_series(
                args.ffmpeg,
                row.src,
                first_time,
                len(group),
                row.source_fps,
                frames_dir,
                f"c{row.n}_verify_s{sample_index}_g{group_index}",
                HIGH_DIM,
                args.ffmpeg_timeout_seconds,
            )
            for offset_frame, image in zip(group, decoded_group):
                verification_frames[(offset_frame, sample_index)] = image

    verified = []
    for offset_frame in sorted(fine):
        delta = offset_frame / row.source_fps
        sample_scores = []
        feature_support = 0
        for i, timeline_time in enumerate(times):
            source = verification_frames.get((offset_frame, i))
            if source is None:
                sample_scores.append(0.0)
                continue
            feature, inliers, ratio = feature_score(source, reference_high[i])
            correlation = max(0.0, min(1.0, (cheap_score(source, reference_high[i]) + 1.0) / 2.0))
            score = correlation * 0.72 + feature * 0.28
            sample_scores.append(score)
            if inliers >= 12 and ratio >= 0.25:
                feature_support += 1
        ordered = sorted(sample_scores)
        median = ordered[len(ordered) // 2]
        combined = median * 0.8 + min(ordered) * 0.2
        verified.append((offset_frame, combined, feature_support, sample_scores))

    verified.sort(key=lambda item: item[1], reverse=True)
    provisional = verified[0]
    exact_radius = 12 if provisional[1] < 0.60 else 2
    exact_verified = []
    for offset_frame in range(max(lo, provisional[0] - exact_radius), min(hi, provisional[0] + exact_radius) + 1):
        sample_scores = []
        feature_support = 0
        for i, timeline_time in enumerate(times):
            source_time = row.in_point + (timeline_time - row.start) + offset_frame / row.source_fps
            source = export_frame(args.ffmpeg, row.src, source_time, frames_dir / f"c{row.n}_exact.png", HIGH_DIM, args.ffmpeg_timeout_seconds)
            feature, inliers, ratio = feature_score(source, reference_high[i])
            correlation = max(0.0, min(1.0, (cheap_score(source, reference_high[i]) + 1.0) / 2.0))
            sample_scores.append(correlation * 0.72 + feature * 0.28)
            if inliers >= 12 and ratio >= 0.25:
                feature_support += 1
        ordered = sorted(sample_scores)
        exact_verified.append((offset_frame, ordered[len(ordered) // 2] * 0.8 + min(ordered) * 0.2, feature_support, sample_scores))
    verified.extend(exact_verified)
    verified.sort(key=lambda item: item[1], reverse=True)
    best = verified[0]
    runner = next((entry for entry in verified[1:] if abs(entry[0] - best[0]) > SHORTLIST_RADIUS_FRAMES), verified[min(1, len(verified) - 1)])
    quality, color = quality_label(best[1], runner[1], best[2], len(times))
    new_in, new_out, applied_offset_seconds = phase_guarded_visual_range(
        row.in_point,
        row.out_point,
        best[0],
        row.source_fps,
    )
    return {
        "status": "MATCHED", "offsetFrames": best[0], "offsetSeconds": best[0] / row.source_fps,
        "appliedOffsetSeconds": applied_offset_seconds,
        "phaseGuardSeconds": PREMIERE_TICK_SECONDS,
        "newIn": new_in,
        "newOut": new_out,
        "score": best[1], "runnerUp": runner[1], "support": best[2], "samples": len(times),
        "quality": quality, "color": color, "lo": lo, "hi": hi,
        "coarseCandidates": len(coarse), "verifiedCandidates": len(verified),
    }


def main() -> int:
    started = time.perf_counter()
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--ffmpeg", required=True)
    parser.add_argument("--frames-dir", required=True)
    parser.add_argument("--records", required=True)
    parser.add_argument("--report", required=True)
    parser.add_argument("--timing-report")
    parser.add_argument("--search-mode", choices=["3_SECONDS", "10_SECONDS", "ENTIRE_SOURCE", "CUSTOM"], default="3_SECONDS")
    parser.add_argument("--custom-seconds", type=float, default=3.0)
    parser.add_argument(
        "--scope",
        choices=["SELECTED", "ALL_TARGET", "ALL_V1"],
        default="SELECTED",
        help="Process selected target clips or all target clips. ALL_V1 is a legacy alias.",
    )
    parser.add_argument("--clip-no", type=int)
    parser.add_argument("--ffmpeg-timeout-seconds", type=float, default=600.0)
    parser.add_argument("--keep-frames", choices=["true", "false"], default="false")
    args = parser.parse_args()

    rows = parse_manifest(Path(args.manifest))
    if args.clip_no is not None:
        rows = [row for row in rows if row.n == args.clip_no]
    frames_dir = Path(args.frames_dir)
    records_path = Path(args.records)
    report_path = Path(args.report)
    results = []
    with records_path.open("w", encoding="utf-8", newline="") as records:
        for row in rows:
            if args.scope == "SELECTED" and not row.selected:
                continue
            try:
                result = analyze(row, args, frames_dir)
            except Exception as exc:
                result = {"status": "NO_MATCH", "reason": clean(exc)}
            result.update({"clipNo": row.n, "clipRef": row.ref, "start": row.start, "duration": row.duration})
            results.append(result)
            if result["status"] == "MATCHED":
                title = f"Slip Match: {result['quality']}"
                method = result.get("matchMethod", "VISUAL_ANALYSIS")
                note = (
                    f"{MARKER_TAG} | Applied slip: {result['offsetFrames']:+d} frames "
                    f"({result['offsetSeconds']:+.3f} s) | Match quality: {result['quality']} | "
                    f"Best score: {result['score']:.3f} | Next-best: {result['runnerUp']:.3f} | "
                    f"Verified samples: {result['support']}/{result['samples']} | Method: {method} | "
                    f"Search: {args.search_mode}"
                )
                fields = ["SLIP", row.ref, result["newIn"], result["newOut"], row.start, row.duration,
                          result["offsetFrames"], result["appliedOffsetSeconds"], result["quality"], result["score"],
                          result["runnerUp"], result["support"], result["samples"], result["color"], title, note]
            else:
                title = "Slip Match: No reliable match"
                note = f"{MARKER_TAG} | Clip was not changed | Reason: {result.get('reason', 'low confidence')} | Search: {args.search_mode}"
                fields = ["MARKER", row.ref, row.in_point, row.out_point, row.start, row.duration,
                          0, 0, "Low", 0, 0, 0, len(times := anchors(row)), 1, title, note]
            records.write("\t".join(clean(value) for value in fields) + "\n")
            records.flush()
            if args.keep_frames != "true":
                for frame_path in frames_dir.glob(f"c{row.n}_*"):
                    frame_path.unlink(missing_ok=True)

    elapsed_seconds = time.perf_counter() - started
    report_path.write_text(
        json.dumps(
            {
                "searchMode": args.search_mode,
                "scope": args.scope,
                "timing": {"processingSeconds": elapsed_seconds},
                "results": results,
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    timing_report = Path(args.timing_report).resolve() if args.timing_report else None
    if timing_report is not None:
        try:
            update_processing_time_report(timing_report, args.clip_no, results, elapsed_seconds)
        except (OSError, ValueError) as exc:
            warn_processing_time_report(timing_report, exc)
    print(
        f"SUMMARY|clips={len(results)}|records={records_path}|report={report_path}"
        f"|timingReport={timing_report or ''}|seconds={elapsed_seconds:.3f}"
    )
    return 0


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