"""
Debug utilities for BADAS inference.

Provides tools for saving inference windows and predictions for analysis.
"""

import cv2
import json
import time
import numpy as np
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional


@dataclass
class DebugConfig:
    """Configuration for debug mode - saves inference windows as videos."""

    enabled: bool = False
    output_dir: str = "debug_windows"
    video_fps: int = 8
    video_codec: str = "libx264"
    video_extension: str = ".mp4"
    save_json: bool = True
    include_timestamp_in_filename: bool = True

    @classmethod
    def disabled(cls) -> 'DebugConfig':
        return cls(enabled=False)

    @classmethod
    def basic(cls, output_dir: str = "debug_windows") -> 'DebugConfig':
        return cls(enabled=True, output_dir=output_dir)

    @classmethod
    def full(cls, output_dir: str = "debug_windows") -> 'DebugConfig':
        return cls(
            enabled=True,
            output_dir=output_dir,
            save_json=True,
            include_timestamp_in_filename=True
        )


class DebugWindowSaver:
    """Saves inference windows as videos for debugging."""

    def __init__(self, config: DebugConfig, img_size: int, model_fps: int, quiet: bool = False):
        self.config = config
        self.img_size = img_size
        self.model_fps = model_fps
        self.quiet = quiet
        self.output_dir = Path(config.output_dir)
        self.results: List[Dict[str, Any]] = []
        self.video_source: str = ""

        if config.enabled:
            self._setup_output_dir()

    def _log(self, msg: str):
        if not self.quiet:
            print(msg)

    def _setup_output_dir(self):
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self._log(f"   Debug output directory: {self.output_dir}")

    def set_video_source(self, video_path: str):
        self.video_source = str(video_path)

    def save_window(
        self,
        frames: List[np.ndarray],
        window_index: int,
        frame_index: int,
        timestamp: float,
        prediction: Dict[str, Any]
    ) -> Optional[str]:
        if not self.config.enabled:
            return None

        ext = ".mp4"

        if self.config.include_timestamp_in_filename:
            filename = f"window_{window_index:04d}_t{timestamp:.2f}s_p{prediction['probability']:.3f}{ext}"
        else:
            filename = f"window_{window_index:04d}{ext}"

        filepath = self.output_dir / filename
        self._write_video(frames, filepath)

        result_entry = {
            "window_index": window_index,
            "frame_index": frame_index,
            "timestamp": round(timestamp, 3),
            "video_file": filename,
            "probability": prediction['probability'],
            "raw_probability": prediction['raw_probability'],
            "risk_level": prediction['risk_level'],
            "latency_ms": prediction['latency_ms']
        }
        self.results.append(result_entry)

        return str(filepath)

    def _write_video(self, frames: List[np.ndarray], filepath: Path):
        import subprocess

        height, width = frames[0].shape[:2]
        filepath = filepath.with_suffix('.mp4')

        cmd = [
            'ffmpeg', '-y',
            '-f', 'rawvideo',
            '-vcodec', 'rawvideo',
            '-s', f'{width}x{height}',
            '-pix_fmt', 'rgb24',
            '-r', str(self.config.video_fps),
            '-i', '-',
            '-c:v', 'libx264',
            '-pix_fmt', 'yuv420p',
            '-preset', 'fast',
            '-crf', '23',
            str(filepath)
        ]

        try:
            proc = subprocess.Popen(
                cmd,
                stdin=subprocess.PIPE,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL
            )

            for frame in frames:
                proc.stdin.write(frame.tobytes())

            proc.stdin.close()
            proc.wait()

        except FileNotFoundError:
            self._log(f"   ffmpeg not found, falling back to OpenCV")
            self._write_video_opencv(frames, filepath)

    def _write_video_opencv(self, frames: List[np.ndarray], filepath: Path):
        height, width = frames[0].shape[:2]
        filepath = filepath.with_suffix('.avi')

        fourcc = cv2.VideoWriter_fourcc(*'MJPG')
        writer = cv2.VideoWriter(str(filepath), fourcc, self.config.video_fps, (width, height))

        if not writer.isOpened():
            self._log(f"   Failed to create video writer for {filepath}")
            return

        for frame in frames:
            bgr_frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            writer.write(bgr_frame)

        writer.release()

    def save_json(self) -> Optional[str]:
        if not self.config.enabled or not self.config.save_json or not self.results:
            return None

        json_path = self.output_dir / "predictions.json"

        probabilities = [r['probability'] for r in self.results]
        raw_probabilities = [r['raw_probability'] for r in self.results]
        latencies = [r['latency_ms'] for r in self.results]

        output = {
            "metadata": {
                "source_video": self.video_source,
                "total_windows": len(self.results),
                "output_directory": str(self.output_dir),
                "created_at": time.strftime("%Y-%m-%d %H:%M:%S")
            },
            "summary": {
                "max_probability": round(max(probabilities), 4),
                "min_probability": round(min(probabilities), 4),
                "mean_probability": round(float(np.mean(probabilities)), 4),
                "max_raw_probability": round(max(raw_probabilities), 4),
                "mean_raw_probability": round(float(np.mean(raw_probabilities)), 4),
                "mean_latency_ms": round(float(np.mean(latencies)), 1),
                "high_risk_windows": sum(1 for p in probabilities if p > 0.7),
                "critical_risk_windows": sum(1 for p in probabilities if p > 0.85)
            },
            "predictions": self.results
        }

        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(output, f, indent=2, ensure_ascii=False)

        self._log(f"   Saved predictions JSON: {json_path}")
        return str(json_path)

    def reset(self):
        self.results.clear()
        self.video_source = ""


__all__ = [
    'DebugConfig',
    'DebugWindowSaver',
]
