import os

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

from vision import Classifier
from utils import box_iou, nms


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 _select_from_list(frames, n):
    if not frames:
        return []
    indices = _sample_indices(len(frames), n)
    return [frames[i] for i in indices]


def _extract_with_cv2(video_path, n):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError("Could not open video file.")

    frames = []
    try:
        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if total > 0:
            for idx in _sample_indices(total, n):
                cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
                ok, frame = cap.read()
                if not ok:
                    continue
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                frames.append(Image.fromarray(frame))
            return frames

        all_frames = []
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            all_frames.append(Image.fromarray(frame))
        return _select_from_list(all_frames, n)
    finally:
        cap.release()


def _get_imageio_count(reader):
    try:
        return reader.count_frames()
    except Exception:
        pass
    try:
        meta = reader.get_meta_data()
        count = meta.get("nframes")
        if count and count != float("inf"):
            return int(count)
    except Exception:
        pass
    return None


def _extract_with_imageio(video_path, n):
    reader = imageio.get_reader(video_path)
    try:
        total = _get_imageio_count(reader)
        if total:
            frames = []
            for idx in _sample_indices(total, n):
                frame = reader.get_data(idx)
                frames.append(Image.fromarray(frame))
            return frames

        all_frames = [Image.fromarray(frame) for frame in reader]
        return _select_from_list(all_frames, n)
    finally:
        reader.close()


def split_video(video_path, n=8):
    if not video_path or not os.path.exists(video_path):
        return []
    # Prefer OpenCV for random access; use imageio as a fallback path if needed.
    try:
        return _extract_with_cv2(video_path, n)
    except Exception:
        return _extract_with_imageio(video_path, n)


model = Classifier(format="onnx")


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):
    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)
    return img


def infer(video_file):
    video_path = _resolve_video_path(video_file)
    frames = split_video(video_path, n=8)
    if not frames:
        return []
    n_frames = len(frames)
    boxes = np.zeros((0, 5), dtype=np.float64)
    frame_preds = []

    for frame in frames:
        bbox = np.asarray(model(frame), dtype=np.float64)
        frame_preds.append(bbox)
        if bbox.size > 0:
            boxes = np.vstack([boxes, bbox])

    if boxes.size == 0:
        return []

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

    # Keep main boxes that appear in enough frames.
    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 // 2
    kept_main = main_bboxes[keep_main] if np.any(keep_main) else np.zeros((0, 5), dtype=np.float64)

    if kept_main.size == 0:
        return []

    outputs = []
    for main_box in 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]))
                outputs.append(_draw_detections(frame, bbox[match_idx : match_idx + 1]))
                break

    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", 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()