import os
import re
import json
import shutil
import subprocess
from typing import Dict, Any, Optional, Callable

def check_ffmpeg_installed() -> bool:
    """Check if ffmpeg is available on the system PATH."""
    return shutil.which("ffmpeg") is not None

def check_ffprobe_installed() -> bool:
    """Check if ffprobe is available on the system PATH."""
    return shutil.which("ffprobe") is not None

def get_video_info(video_path: str) -> Dict[str, Any]:
    """
    Get video metadata using ffprobe or ffmpeg.
    Returns dictionary with fps, width, height, duration, frame_count, has_audio.
    """
    info = {
        "fps": 30.0,
        "width": 1920,
        "height": 1080,
        "duration": 0.0,
        "frame_count": 0,
        "has_audio": False
    }
    
    if check_ffprobe_installed():
        cmd = [
            "ffprobe",
            "-v", "quiet",
            "-print_format", "json",
            "-show_format",
            "-show_streams",
            video_path
        ]
        try:
            result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
            data = json.loads(result.stdout)
            
            for stream in data.get("streams", []):
                if stream.get("codec_type") == "video" and "r_frame_rate" in stream:
                    # Evaluate fps string like "30/1" or "30000/1001"
                    r_fps = stream["r_frame_rate"]
                    if "/" in r_fps:
                        num, den = r_fps.split("/")
                        if float(den) > 0:
                            info["fps"] = float(num) / float(den)
                    else:
                        info["fps"] = float(r_fps)
                        
                    info["width"] = int(stream.get("width", 1920))
                    info["height"] = int(stream.get("height", 1080))
                    
                    if "nb_frames" in stream and stream["nb_frames"].isdigit():
                        info["frame_count"] = int(stream["nb_frames"])
                elif stream.get("codec_type") == "audio":
                    info["has_audio"] = True
                    
            if "format" in data and "duration" in data["format"]:
                info["duration"] = float(data["format"]["duration"])
                if info["frame_count"] == 0 and info["fps"] > 0:
                    info["frame_count"] = int(info["duration"] * info["fps"])
            return info
        except Exception:
            pass

    # Fallback parsing via ffmpeg stderr
    cmd = ["ffmpeg", "-i", video_path]
    res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    stderr = res.stderr

    fps_match = re.search(r"(\d+(?:\.\d+)?)\s+fps", stderr)
    if fps_match:
        info["fps"] = float(fps_match.group(1))

    res_match = re.search(r"(\d{2,5})x(\d{2,5})", stderr)
    if res_match:
        info["width"] = int(res_match.group(1))
        info["height"] = int(res_match.group(2))

    dur_match = re.search(r"Duration:\s*(\d+):(\d+):(\d+(?:\.\d+)?)", stderr)
    if dur_match:
        h, m, s = float(dur_match.group(1)), float(dur_match.group(2)), float(dur_match.group(3))
        info["duration"] = h * 3600 + m * 60 + s
        info["frame_count"] = int(info["duration"] * info["fps"])

    if "Audio:" in stderr:
        info["has_audio"] = True

    return info

def extract_frames_lossless(
    video_path: str,
    output_dir: str,
    target_fps: Optional[float] = None,
    progress_callback: Optional[Callable[[str], None]] = None
) -> int:
    """
    Extract frames from a video into PNG images (100% Lossless) inside output_dir.
    Returns total number of extracted frames.
    """
    os.makedirs(output_dir, exist_ok=True)
    out_pattern = os.path.join(output_dir, "frame_%06d.png")

    cmd = ["ffmpeg", "-y", "-i", video_path]

    if target_fps and target_fps > 0:
        cmd.extend(["-vf", f"fps={target_fps}"])

    # Lossless PNG output codec (-c:v png -pred mixed)
    cmd.extend(["-c:v", "png", "-pred", "mixed", out_pattern])

    if progress_callback:
        progress_callback("Extracting video frames in lossless PNG format...")

    process = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if process.returncode != 0:
        raise RuntimeError(f"FFmpeg frame extraction failed: {process.stderr}")

    frames = [f for f in os.listdir(output_dir) if f.startswith("frame_") and f.endswith(".png")]
    frames.sort()
    
    if progress_callback:
        progress_callback(f"Successfully extracted {len(frames)} lossless frames.")

    return len(frames)

def extract_single_frame(video_path: str, output_path: str, time_offset: float = 0.5) -> str:
    """Grab one representative frame from a video (for preview purposes)."""
    cmd = ["ffmpeg", "-y", "-ss", str(time_offset), "-i", video_path, "-frames:v", "1", output_path]
    res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if res.returncode != 0 or not os.path.exists(output_path):
        raise RuntimeError(f"FFmpeg single-frame extraction failed: {res.stderr}")
    return output_path

def extract_audio(video_path: str, output_audio_path: str) -> bool:
    """Extract audio stream from video file if present."""
    cmd = ["ffmpeg", "-y", "-i", video_path, "-vn", "-acodec", "aac", "-b:a", "192k", output_audio_path]
    res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return res.returncode == 0 and os.path.exists(output_audio_path) and os.path.getsize(output_audio_path) > 0

def reassemble_video(
    frames_dir: str,
    output_video_path: str,
    fps: float = 30.0,
    audio_path: Optional[str] = None,
    export_format: str = "mov_transparent",
    bg_color: Optional[str] = None,
    progress_callback: Optional[Callable[[str], None]] = None
) -> str:
    """
    Reassemble processed frame sequence back into a video file or frame folder.
    
    export_format options:
    - 'mov_transparent': QuickTime MOV with ProRes 4444 + Alpha Transparency
    - 'webm_transparent': WebM with VP9 + Alpha Transparency
    - 'mp4_solid': MP4 (H.264) with solid background color (e.g. Green screen #00FF00)
    - 'png_sequence': Copies/Leaves processed PNG frames in output directory
    """
    if export_format == "png_sequence":
        out_dir = output_video_path
        os.makedirs(out_dir, exist_ok=True)
        for fname in os.listdir(frames_dir):
            if fname.endswith(".png"):
                shutil.copy2(os.path.join(frames_dir, fname), os.path.join(out_dir, fname))
        if progress_callback:
            progress_callback(f"PNG sequence exported to {out_dir}")
        return out_dir

    input_pattern = os.path.join(frames_dir, "frame_%06d.png")
    
    # Ensure parent output directory exists
    os.makedirs(os.path.dirname(os.path.abspath(output_video_path)), exist_ok=True)

    cmd = ["ffmpeg", "-y", "-framerate", str(fps), "-i", input_pattern]

    if audio_path and os.path.exists(audio_path):
        cmd.extend(["-i", audio_path, "-c:a", "aac", "-shortest"])

    if export_format == "mov_transparent":
        # ProRes 4444 with Alpha Channel (Lossless/visually lossless transparent video)
        cmd.extend([
            "-c:v", "prores_ks",
            "-profile:v", "4",
            "-pix_fmt", "yuva444p10le",
            output_video_path
        ])
    elif export_format == "webm_transparent":
        # VP9 WebM with Alpha Channel
        cmd.extend([
            "-c:v", "libvpx-vp9",
            "-pix_fmt", "yuva420p",
            "-b:v", "4M",
            output_video_path
        ])
    elif export_format == "mp4_solid":
        # Standard MP4 H.264
        cmd.extend([
            "-c:v", "libx264",
            "-pix_fmt", "yuv420p",
            "-crf", "18",
            output_video_path
        ])
    else:
        # Default MP4 fallback
        cmd.extend([
            "-c:v", "libx264",
            "-pix_fmt", "yuv420p",
            output_video_path
        ])

    if progress_callback:
        progress_callback("Reassembling frames into output video format...")

    res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if res.returncode != 0:
        # Fallback to standard libx264 if prores_ks or libvpx-vp9 isn't available
        if export_format in ["mov_transparent", "webm_transparent"]:
            if progress_callback:
                progress_callback(f"Codec warning: trying fallback MP4 export for {output_video_path}...")
            fallback_cmd = [
                "ffmpeg", "-y", "-framerate", str(fps), "-i", input_pattern
            ]
            if audio_path and os.path.exists(audio_path):
                fallback_cmd.extend(["-i", audio_path, "-c:a", "aac", "-shortest"])
            fallback_cmd.extend(["-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "18", output_video_path])
            res_fb = subprocess.run(fallback_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            if res_fb.returncode != 0:
                raise RuntimeError(f"FFmpeg video reassembly failed: {res_fb.stderr}")
        else:
            raise RuntimeError(f"FFmpeg video reassembly failed: {res.stderr}")

    if progress_callback:
        progress_callback(f"Video created successfully: {output_video_path}")

    return output_video_path