import logging
import os
import shutil
import subprocess
import tempfile
import time
from collections import deque
from contextlib import contextmanager

import cv2
import gradio as gr
import numpy as np
from PIL import Image, ImageDraw

from vision import Classifier
from utils import box_iou, nms


LOGGER = logging.getLogger(__name__)


DEFAULT_SPLIT_CFG = {
    "n_samples": 30,
    "max_w": 400,
    "crop_y": (0.25, 0.90),
    "dx_threshold_px": 1.5,
    "min_inlier_ratio": 0.20,
    "min_stable_frames": 2,
    "smooth_window": 2,
    "orb_nfeatures": 800,
    "orb_fast_threshold": 12,
    "min_matches": 25,
    "keep_ratio": 0.4,
    "jump_meanabs_threshold": 18.0,
    "progress_every": 0,
}


def _sample_indices(total, n):
    if total <= 0:
        return []
    if total <= n:
        return list(range(total))
    return np.linspace(0, total - 1, n).astype(int).tolist()


def _format_idx_list(indices, max_items=40):
    if not indices:
        return "[]"
    values = [int(i) for i in indices]
    if len(values) <= max_items:
        return str(values)
    head = values[: max_items // 2]
    tail = values[-(max_items // 2) :]
    return f"{head} ... {tail} (len={len(values)})"


def _parse_fraction(value):
    if not value:
        return None
    txt = str(value).strip()
    if not txt or txt == "0/0":
        return None
    if "/" in txt:
        num, den = txt.split("/", 1)
        try:
            den_f = float(den)
            if den_f == 0:
                return None
            return float(num) / den_f
        except Exception:
            return None
    try:
        return float(txt)
    except Exception:
        return None


def _probe_total_frames_ffprobe(video_path):
    ffprobe = shutil.which("ffprobe")
    if ffprobe is None:
        return None

    # Try direct frame count first.
    cmd = [
        ffprobe,
        "-v",
        "error",
        "-select_streams",
        "v:0",
        "-show_entries",
        "stream=nb_frames",
        "-of",
        "default=noprint_wrappers=1:nokey=1",
        video_path,
    ]
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False)
    if proc.returncode == 0:
        raw = proc.stdout.strip()
        if raw.isdigit():
            val = int(raw)
            if val > 0:
                return val

    # Fallback: estimate from duration * avg frame rate.
    cmd = [
        ffprobe,
        "-v",
        "error",
        "-select_streams",
        "v:0",
        "-show_entries",
        "stream=avg_frame_rate,duration",
        "-of",
        "default=noprint_wrappers=1:nokey=1",
        video_path,
    ]
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False)
    if proc.returncode != 0:
        return None

    lines = [line.strip() for line in proc.stdout.splitlines() if line.strip()]
    if len(lines) < 2:
        return None

    fps = _parse_fraction(lines[0])
    duration = _parse_fraction(lines[1])
    if fps is None or duration is None:
        return None

    estimate = int(round(fps * duration))
    return estimate if estimate > 0 else None


def _extract_bgr_with_ffmpeg(video_path, n):
    ffmpeg = shutil.which("ffmpeg")
    if ffmpeg is None:
        raise RuntimeError("ffmpeg is not available")

    total = _probe_total_frames_ffprobe(video_path)
    if total is None or total <= 0:
        raise RuntimeError("ffprobe could not determine total frame count")

    indices = _sample_indices(total, int(n))
    if not indices:
        return []

    LOGGER.info(
        "Frame extraction | video=%s total_frames=%d n_samples=%d sampled_indices=%s",
        os.path.basename(video_path),
        total,
        len(indices),
        _format_idx_list(indices),
    )

    select_expr = "+".join(f"eq(n\\,{int(i)})" for i in indices)
    vf = f"select={select_expr}"

    with tempfile.TemporaryDirectory(prefix="ffmpeg_frames_") as tmpdir:
        pattern = os.path.join(tmpdir, "frame_%06d.jpg")
        cmd = [
            ffmpeg,
            "-hide_banner",
            "-loglevel",
            "error",
            "-i",
            video_path,
            "-vf",
            vf,
            "-vsync",
            "vfr",
            "-q:v",
            "2",
            pattern,
        ]
        proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False)
        if proc.returncode != 0:
            raise RuntimeError(proc.stderr.strip() or "ffmpeg extraction failed")

        frames = []
        for name in sorted(os.listdir(tmpdir)):
            if not name.lower().endswith(".jpg"):
                continue
            frame = cv2.imread(os.path.join(tmpdir, name), cv2.IMREAD_COLOR)
            if frame is not None:
                frames.append(frame)
        LOGGER.info(
            "Frame extraction done | video=%s extracted=%d requested=%d",
            os.path.basename(video_path),
            len(frames),
            len(indices),
        )
        return frames


def _extract_with_ffmpeg(video_path, n):
    frames = _extract_bgr_with_ffmpeg(video_path, n)
    return [Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) for frame in frames]


def split_video(video_path, n=8):
    if not video_path or not os.path.exists(video_path):
        return []
    return _extract_with_ffmpeg(video_path, n)


@contextmanager
def timer(name, stats):
    t0 = time.perf_counter()
    yield
    stats[name] = stats.get(name, 0.0) + (time.perf_counter() - t0)


def _iter_sampled_frames(video_path, n_samples):
    frames = _extract_bgr_with_ffmpeg(video_path, int(n_samples))
    for out_idx, frame in enumerate(frames):
        yield out_idx, frame


def iter_frames(video_path, n_samples, max_w, crop_y):
    for out_idx, frame in _iter_sampled_frames(video_path, n_samples):
        proc = frame
        if max_w > 0 and proc.shape[1] != max_w:
            scale = max_w / float(proc.shape[1])
            proc = cv2.resize(
                proc,
                (max_w, int(proc.shape[0] * scale)),
                interpolation=cv2.INTER_AREA,
            )

        if crop_y is not None:
            h = proc.shape[0]
            y0 = int(max(0.0, min(1.0, float(crop_y[0]))) * h)
            y1 = int(max(0.0, min(1.0, float(crop_y[1]))) * h)
            if y1 > y0:
                proc = proc[y0:y1, :]

        yield out_idx, proc


def quick_jump_score(prev_gray, gray, small_w=160):
    h, w = prev_gray.shape[:2]
    if w > small_w:
        scale = small_w / float(w)
        prev_s = cv2.resize(prev_gray, (small_w, int(h * scale)), interpolation=cv2.INTER_AREA)
        gray_s = cv2.resize(gray, (small_w, int(h * scale)), interpolation=cv2.INTER_AREA)
    else:
        prev_s = prev_gray
        gray_s = gray

    diff = cv2.absdiff(prev_s, gray_s)
    return float(np.mean(diff))


def estimate_dx_orb_affine(prev_gray, gray, orb, bf, min_matches, keep_ratio, timing_pair):
    with timer("orb_detect_compute", timing_pair):
        kp1, des1 = orb.detectAndCompute(prev_gray, None)
        kp2, des2 = orb.detectAndCompute(gray, None)

    if des1 is None or des2 is None or len(kp1) < 8 or len(kp2) < 8:
        return None

    with timer("bf_match", timing_pair):
        matches = bf.match(des1, des2)

    if len(matches) < min_matches:
        return None

    with timer("match_sort_filter", timing_pair):
        matches = sorted(matches, key=lambda m: m.distance)
        keep_n = max(8, int(len(matches) * keep_ratio))
        matches = matches[:keep_n]

        pts1 = np.float32([kp1[m.queryIdx].pt for m in matches])
        pts2 = np.float32([kp2[m.trainIdx].pt for m in matches])

    with timer("ransac_affine", timing_pair):
        M, inliers = cv2.estimateAffinePartial2D(
            pts1,
            pts2,
            method=cv2.RANSAC,
            ransacReprojThreshold=3.0,
            maxIters=1500,
            confidence=0.99,
        )

    if M is None:
        return None

    dx = float(M[0, 2])
    dy = float(M[1, 2])
    inlier_ratio = float(np.mean(inliers)) if inliers is not None else 0.0

    return {
        "dx": dx,
        "dy": dy,
        "score_dx": float(abs(dx)),
        "score_px": float(np.hypot(dx, dy)),
        "inlier_ratio": inlier_ratio,
        "matches": len(matches),
        "M": M,
    }


def split_video_into_stable_segments_fast(
    video_path,
    n_samples=30,
    max_w=400,
    crop_y=(0.25, 0.90),
    dx_threshold_px=1.5,
    min_inlier_ratio=0.20,
    min_stable_frames=2,
    smooth_window=2,
    orb_nfeatures=800,
    orb_fast_threshold=12,
    min_matches=25,
    keep_ratio=0.4,
    jump_meanabs_threshold=18.0,
    progress_every=200,
):
    timing_total = {}
    timing_pair = {}

    with timer("setup", timing_total):
        orb = cv2.ORB_create(nfeatures=orb_nfeatures, fastThreshold=orb_fast_threshold)
        bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)

    metrics = []
    prev_gray = None
    frame_count = 0

    with timer("loop_total", timing_total):
        for _, frame in iter_frames(video_path, n_samples=n_samples, max_w=max_w, crop_y=crop_y):
            frame_count += 1

            with timer("to_gray", timing_total):
                gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

            if prev_gray is not None:
                with timer("quick_jump", timing_total):
                    q = quick_jump_score(prev_gray, gray)

                if q >= jump_meanabs_threshold:
                    metrics.append(
                        {
                            "dx": np.nan,
                            "dy": np.nan,
                            "score_dx": 1e9,
                            "score_px": 1e9,
                            "inlier_ratio": 0.0,
                            "matches": 0,
                            "M": None,
                            "quick_jump": q,
                        }
                    )
                else:
                    m = estimate_dx_orb_affine(
                        prev_gray,
                        gray,
                        orb=orb,
                        bf=bf,
                        min_matches=min_matches,
                        keep_ratio=keep_ratio,
                        timing_pair=timing_pair,
                    )
                    if m is None:
                        metrics.append(
                            {
                                "dx": np.nan,
                                "dy": np.nan,
                                "score_dx": 1e9,
                                "score_px": 1e9,
                                "inlier_ratio": 0.0,
                                "matches": 0,
                                "M": None,
                                "quick_jump": q,
                            }
                        )
                    else:
                        m["quick_jump"] = q
                        metrics.append(m)

                if progress_every and (len(metrics) % progress_every == 0):
                    print(f"processed pairs: {len(metrics)}")

            prev_gray = gray

    if frame_count < 2:
        return [], metrics, [], {"total": timing_total, "per_pair": timing_pair}

    with timer("post_smooth", timing_total):
        raw_dx = [m["score_dx"] for m in metrics]
        raw_inlier = [m["inlier_ratio"] for m in metrics]

        smoothed_dx = []
        q = deque(maxlen=max(1, int(smooth_window)))
        for v in raw_dx:
            if not np.isfinite(v):
                q.clear()
                smoothed_dx.append(np.nan)
            else:
                q.append(v)
                smoothed_dx.append(float(np.mean(q)))

    with timer("post_segments", timing_total):
        min_len = max(1, int(min_stable_frames))

        stable_flags = []
        for dx_s, r in zip(smoothed_dx, raw_inlier):
            if not np.isfinite(dx_s):
                stable_flags.append(False)
            else:
                stable_flags.append((dx_s < dx_threshold_px) and (r >= min_inlier_ratio))

        segments = []
        start = None
        for i, is_stable in enumerate(stable_flags):
            if is_stable and start is None:
                start = i
            if (not is_stable) and start is not None:
                end = i
                if (end - start) >= min_len:
                    segments.append((start, end))
                start = None

        if start is not None:
            end = len(stable_flags)
            if (end - start) >= min_len:
                segments.append((start, end))

    LOGGER.info(
        "Segmentation summary | sampled_frames=%d pair_metrics=%d stable_segments=%d",
        frame_count,
        len(metrics),
        len(segments),
    )
    if segments:
        LOGGER.info("Segment ranges (sample indices) | %s", segments)

    return segments, metrics, smoothed_dx, {"total": timing_total, "per_pair": timing_pair}


def _bgr_to_pil(frame):
    return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))


def extract_segment_frames(video_path, segments, n_samples):
    if not segments:
        LOGGER.info("Segment frame extraction | no segments found")
        return []

    normalized_segments = []
    for start, end in segments:
        s = max(0, int(start))
        e = max(s, int(end))
        normalized_segments.append((s, e))

    normalized_segments.sort(key=lambda x: x[0])
    grouped_frames = [[] for _ in normalized_segments]
    grouped_indices = [[] for _ in normalized_segments]
    segment_idx = 0

    # Detection runs on original sampled frames (no resize / no crop).
    for frame_idx, frame in _iter_sampled_frames(video_path, n_samples=n_samples):
        while segment_idx < len(normalized_segments) and frame_idx > normalized_segments[segment_idx][1]:
            segment_idx += 1

        if segment_idx >= len(normalized_segments):
            break

        seg_start, seg_end = normalized_segments[segment_idx]
        if seg_start <= frame_idx <= seg_end:
            grouped_frames[segment_idx].append(_bgr_to_pil(frame))
            grouped_indices[segment_idx].append(frame_idx)

    LOGGER.info(
        "Segment frame extraction summary | segments=%d n_samples=%d",
        len(normalized_segments),
        n_samples,
    )
    for seg_i, ((seg_start, seg_end), idx_list, frames) in enumerate(
        zip(normalized_segments, grouped_indices, grouped_frames),
        start=1,
    ):
        LOGGER.info(
            "Segment %d | requested_range=[%d,%d] matched_frames=%d matched_indices=%s",
            seg_i,
            seg_start,
            seg_end,
            len(frames),
            _format_idx_list(idx_list),
        )

    return [frames for frames in grouped_frames if frames]


def split_video_stable(video_path, split_cfg=None, fallback_n=30):
    if not video_path or not os.path.exists(video_path):
        return []

    cfg = DEFAULT_SPLIT_CFG.copy()
    if split_cfg:
        cfg.update(split_cfg)

    LOGGER.info("Split config | %s", cfg)

    segments, _, _, _ = split_video_into_stable_segments_fast(video_path, **cfg)
    frame_groups = extract_segment_frames(
        video_path,
        segments,
        n_samples=cfg["n_samples"],
    )

    if frame_groups:
        LOGGER.info(
            "Split result | stable_splits=%d split_frame_counts=%s",
            len(frame_groups),
            [len(group) for group in frame_groups],
        )
        return frame_groups

    LOGGER.info("Split result | no stable segment, using fallback sampling n=%d", fallback_n)
    fallback_frames = split_video(video_path, n=fallback_n)
    LOGGER.info("Fallback frame count | %d", len(fallback_frames))
    return [fallback_frames] if fallback_frames else []


model = Classifier(format="onnx", conf=0.05)


def _resolve_video_path(video_input):
    if not video_input:
        return None
    if isinstance(video_input, str):
        return video_input
    if isinstance(video_input, dict):
        for key in ("name", "path", "data", "video"):
            value = video_input.get(key)
            if isinstance(value, str) and os.path.exists(value):
                return value
    if isinstance(video_input, (list, tuple)):
        for value in video_input:
            if isinstance(value, str) and os.path.exists(value):
                return value
    return None


def _draw_detections(pil_img, preds, subtitle=None):
    img = pil_img.copy()
    draw = ImageDraw.Draw(img)
    width, height = img.size
    color = (255, 80, 0)
    preds = np.asarray(preds)

    for x1, y1, x2, y2, conf in preds:
        x1 = int(max(0.0, min(1.0, float(x1))) * width)
        y1 = int(max(0.0, min(1.0, float(y1))) * height)
        x2 = int(max(0.0, min(1.0, float(x2))) * width)
        y2 = int(max(0.0, min(1.0, float(y2))) * height)
        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
        draw.text((x1 + 4, y1 + 4), f"{conf:.2f}", fill=color)

    draw.text((6, 6), f"detections: {len(preds)}", fill=color)
    if subtitle:
        draw.text((6, 26), subtitle, fill=color)
    return img


def _combine_predictions_per_split(frame_preds):
    n_frames = len(frame_preds)
    if n_frames == 0:
        return np.zeros((0, 5), dtype=np.float64)

    boxes = np.zeros((0, 5), dtype=np.float64)
    for bbox in frame_preds:
        if bbox.size > 0:
            boxes = np.vstack([boxes, bbox])

    if boxes.size == 0:
        return np.zeros((0, 5), dtype=np.float64)

    main_bboxes = np.asarray(nms(boxes), dtype=np.float64)
    if main_bboxes.size == 0:
        return np.zeros((0, 5), dtype=np.float64)

    matches_per_main = np.zeros(len(main_bboxes), dtype=int)
    for bbox in frame_preds:
        if bbox.size == 0:
            continue
        ious = box_iou(bbox[:, :4], main_bboxes[:, :4])
        matches_per_main += (ious > 0).any(axis=1).astype(int)

    keep_main = matches_per_main > n_frames // 4
    if np.any(keep_main):
        return main_bboxes[keep_main]
    return np.zeros((0, 5), dtype=np.float64)


def infer(video_file):
    video_path = _resolve_video_path(video_file)
    LOGGER.info("Inference start | video=%s", video_path)
    split_frames = split_video_stable(video_path)
    if not split_frames:
        LOGGER.info("Inference stop | no frames available")
        return []

    outputs = []
    for split_idx, frames in enumerate(split_frames):
        LOGGER.info("Inference split %d | frames=%d", split_idx + 1, len(frames))
        frame_preds = []
        for frame in frames:
            bbox = np.asarray(model(frame), dtype=np.float64).reshape(-1, 5)
            frame_preds.append(bbox)

        kept_main = _combine_predictions_per_split(frame_preds)
        LOGGER.info(
            "Inference split %d | combined_detections=%d",
            split_idx + 1,
            len(kept_main),
        )
        if kept_main.size == 0:
            continue

        for det_idx, main_box in enumerate(kept_main):
            for frame, bbox in zip(frames, frame_preds):
                if bbox.size == 0:
                    continue
                ious = box_iou(bbox[:, :4], main_box[:4].reshape(1, 4))
                if (ious > 0).any():
                    match_idx = int(np.argmax(ious[0]))
                    subtitle = f"split {split_idx + 1} / det {det_idx + 1}"
                    outputs.append(_draw_detections(frame, bbox[match_idx : match_idx + 1], subtitle=subtitle))
                    break

    LOGGER.info("Inference done | output_images=%d", len(outputs))
    return outputs


with gr.Blocks() as demo:
    gr.Markdown("## Pyronear Wildfire Detection")
    with gr.Row():
        video_in = gr.Video(label="Upload MP4")
    gallery_out = gr.Gallery(label="Wildfire detected (per split)", columns=2, height=360)
    run_btn = gr.Button("Detect")
    run_btn.click(fn=infer, inputs=video_in, outputs=gallery_out)


if __name__ == "__main__":
    demo.launch()