from __future__ import annotations

import shutil
import time
from pathlib import Path
from typing import Dict, List, Optional

import cv2

from app.config import settings
from app.pipeline import BirdCascadePipeline
from app.utils import (
    draw_results,
    now_timestamp_minute,
    safe_stem,
    save_json,
    write_video_xml,
)


def make_unique_result_dir(parent_dir: Path, timestamp: str, original_filename: str) -> tuple[Path, str, str]:
    """
    创建单次视频识别结果目录。

    目录名：
        时间戳_原文件名stem

    例如：
        runs/video/202605181431_DJI_20260423155914_0006_Z/

    返回：
        result_dir, result_base_name, original_stem
    """
    original_stem = safe_stem(original_filename)
    base_name = f"{timestamp}_{original_stem}"

    result_dir = parent_dir / base_name

    if not result_dir.exists():
        result_dir.mkdir(parents=True, exist_ok=True)
        return result_dir, base_name, original_stem

    idx = 1
    while True:
        candidate_base = f"{base_name}_{idx:02d}"
        candidate_dir = parent_dir / candidate_base
        if not candidate_dir.exists():
            candidate_dir.mkdir(parents=True, exist_ok=True)
            return candidate_dir, candidate_base, original_stem
        idx += 1


def iou_xyxy(a: List[float], b: List[float]) -> float:
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b

    x1 = max(ax1, bx1)
    y1 = max(ay1, by1)
    x2 = min(ax2, bx2)
    y2 = min(ay2, by2)

    inter = max(0.0, x2 - x1) * max(0.0, y2 - y1)
    area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1)
    area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1)

    return inter / (area_a + area_b - inter + 1e-7)


class VideoResultSmoother:
    """
    降低视频框闪烁：
    1. 非推理帧复用上一推理结果；
    2. 短暂漏检时保留上一结果若干帧；
    3. 对匹配目标的坐标做指数平滑；
    4. 对同一目标累计物种置信度，避免标签和颜色逐帧跳变。
    """

    def __init__(self, hold_frames: int = 8, alpha: float = 0.65):
        self.hold_frames = hold_frames
        self.alpha = alpha
        self.last_results: List[Dict] = []
        self.last_update_frame: int = -10**9

    @staticmethod
    def _public_results(results: List[Dict]) -> List[Dict]:
        return [
            {
                key: value
                for key, value in item.items()
                if not key.startswith("_smoother_")
            }
            for item in results
        ]

    def update(self, current_results: List[Dict], frame_index: int) -> List[Dict]:
        if len(current_results) == 0:
            if frame_index - self.last_update_frame <= self.hold_frames:
                return self._public_results(self.last_results)
            self.last_results = []
            return []

        if len(self.last_results) == 0:
            self.last_results = []
            for current in current_results:
                item = dict(current)
                class_key = (
                    int(item.get("cls_class_id", -1)),
                    str(item.get("cls_class_name", "") or ""),
                )
                confidence = max(float(item.get("cls_confidence", 0.0)), 0.01)
                item["_smoother_class_votes"] = {class_key: confidence}
                self.last_results.append(item)
            self.last_update_frame = frame_index
            return self._public_results(self.last_results)

        smoothed_results = []
        used_last = set()

        for cur in current_results:
            cur_box = cur["bbox_xyxy"]

            best_idx = -1
            best_iou = 0.0

            for idx, prev in enumerate(self.last_results):
                if idx in used_last:
                    continue

                score = iou_xyxy(cur_box, prev["bbox_xyxy"])
                if score > best_iou:
                    best_iou = score
                    best_idx = idx

            new_item = dict(cur)

            if best_idx >= 0 and best_iou >= 0.1:
                prev = self.last_results[best_idx]
                prev_box = prev["bbox_xyxy"]

                smooth_box = []
                for p, c in zip(prev_box, cur_box):
                    smooth_box.append(round(self.alpha * float(p) + (1 - self.alpha) * float(c), 2))

                new_item["bbox_xyxy"] = smooth_box

                votes = {
                    key: float(value) * 0.90
                    for key, value in prev.get(
                        "_smoother_class_votes",
                        {},
                    ).items()
                }
                current_key = (
                    int(cur.get("cls_class_id", -1)),
                    str(cur.get("cls_class_name", "") or ""),
                )
                current_confidence = max(
                    float(cur.get("cls_confidence", 0.0)),
                    0.01,
                )
                votes[current_key] = votes.get(current_key, 0.0) + current_confidence
                chosen_key = max(votes, key=votes.get)
                new_item["_smoother_class_votes"] = votes
                new_item["cls_class_id"] = chosen_key[0]
                new_item["cls_class_name"] = chosen_key[1]
                if chosen_key != current_key:
                    new_item["show_cls_name"] = bool(
                        prev.get("show_cls_name", False)
                    )
                used_last.add(best_idx)
            else:
                class_key = (
                    int(cur.get("cls_class_id", -1)),
                    str(cur.get("cls_class_name", "") or ""),
                )
                confidence = max(float(cur.get("cls_confidence", 0.0)), 0.01)
                new_item["_smoother_class_votes"] = {class_key: confidence}

            smoothed_results.append(new_item)

        self.last_results = smoothed_results
        self.last_update_frame = frame_index

        return self._public_results(smoothed_results)

    def get_for_skipped_frame(self, frame_index: int) -> List[Dict]:
        if frame_index - self.last_update_frame <= self.hold_frames:
            return self._public_results(self.last_results)
        return []


def process_video_file(
    video_path: str,
    pipeline: BirdCascadePipeline,
    original_filename: str,
    sample_interval: Optional[int] = None,
    save_annotated_video: Optional[bool] = None,
) -> Dict:
    """
    视频结果自动保存到：
        runs/video/时间戳_原视频名/

    目录内部生成：
        原视频名.mp4
        时间戳_原视频名.json
        时间戳_原视频名.xml
        原视频名_source.原后缀
        frames/

    示例：
        runs/video/202605181431_DJI_20260423155914_0006_Z/
        ├── DJI_20260423155914_0006_Z.mp4
        ├── 202605181431_DJI_20260423155914_0006_Z.json
        ├── 202605181431_DJI_20260423155914_0006_Z.xml
        ├── DJI_20260423155914_0006_Z_source.MP4
        └── frames/
    """
    sample_interval = settings.video_sample_interval if sample_interval is None else int(sample_interval)
    save_annotated_video = settings.save_annotated_video if save_annotated_video is None else bool(save_annotated_video)

    video_root = Path(settings.output_dir) / settings.video_output_subdir
    video_root.mkdir(parents=True, exist_ok=True)

    timestamp = now_timestamp_minute()
    result_dir, result_base_name, original_stem = make_unique_result_dir(
        parent_dir=video_root,
        timestamp=timestamp,
        original_filename=original_filename,
    )

    original_suffix = Path(original_filename).suffix or ".mp4"

    source_video_path = result_dir / f"{original_stem}_source{original_suffix}"

    # 输出视频固定为 mp4，文件名不带时间戳
    annotated_video_path = result_dir / f"{original_stem}.mp4"

    # JSON 和 XML 带时间戳
    json_path = result_dir / f"{result_base_name}.json"
    xml_path = result_dir / f"{result_base_name}.xml"

    frames_dir = result_dir / "frames"
    frames_dir.mkdir(parents=True, exist_ok=True)

    shutil.copy2(video_path, source_video_path)

    cap = cv2.VideoCapture(str(source_video_path))
    if not cap.isOpened():
        raise ValueError(f"视频打开失败: {source_video_path}")

    fps = cap.get(cv2.CAP_PROP_FPS)
    if fps is None or fps <= 0:
        fps = 25.0

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)

    writer = None
    if save_annotated_video:
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        writer = cv2.VideoWriter(str(annotated_video_path), fourcc, fps, (width, height))

    smoother = VideoResultSmoother(
        hold_frames=settings.smooth_hold_frames,
        alpha=settings.box_smooth_alpha,
    )

    frame_results = []
    processed_frames = 0
    detected_frames = 0
    frame_index = 0

    start_time = time.time()

    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                break

            should_process = frame_index % sample_interval == 0

            if should_process:
                output = pipeline.predict(frame, return_image=False)
                raw_results = output["results"]
                processed_frames += 1

                if settings.smooth_video_boxes:
                    results = smoother.update(raw_results, frame_index)
                else:
                    results = raw_results

                if len(results) > 0:
                    detected_frames += 1

                frame_results.append(
                    {
                        "frame_index": frame_index,
                        "time_sec": round(frame_index / fps, 3),
                        "num_objects": len(results),
                        "results": results,
                    }
                )
            else:
                if settings.smooth_video_boxes:
                    results = smoother.get_for_skipped_frame(frame_index)
                else:
                    results = []

            annotated_frame = draw_results(frame, results)

            if should_process and len(results) > 0:
                frame_img_path = frames_dir / f"frame_{frame_index:08d}.jpg"
                cv2.imwrite(str(frame_img_path), annotated_frame)

            if writer is not None:
                writer.write(annotated_frame)

            frame_index += 1

    finally:
        cap.release()
        if writer is not None:
            writer.release()

    elapsed = round(time.time() - start_time, 3)

    summary = {
        "success": True,
        "type": "video",
        "original_filename": original_filename,
        "result_dir": str(result_dir),
        "source_video": str(source_video_path),
        "annotated_video": str(annotated_video_path) if annotated_video_path.exists() else None,
        "json": str(json_path),
        "xml": str(xml_path),
        "frames_dir": str(frames_dir),
        "video_info": {
            "width": width,
            "height": height,
            "fps": fps,
            "total_frames": total_frames,
        },
        "sample_interval": sample_interval,
        "smooth_video_boxes": settings.smooth_video_boxes,
        "smooth_hold_frames": settings.smooth_hold_frames,
        "box_smooth_alpha": settings.box_smooth_alpha,
        "processed_frames": processed_frames,
        "detected_frames": detected_frames,
        "elapsed_sec": elapsed,
    }

    json_data = {
        "summary": summary,
        "frames": frame_results,
    }

    save_json(json_path, json_data)

    write_video_xml(
        xml_path=xml_path,
        video_filename=annotated_video_path.name,
        video_path=str(annotated_video_path),
        width=width,
        height=height,
        fps=fps,
        total_frames=total_frames,
        sample_interval=sample_interval,
        frame_results=frame_results,
    )

    return summary