from __future__ import annotations

import json
import re
import subprocess
from pathlib import Path

import numpy as np
from PIL import Image, ImageFilter


def matte_for_frame(image: Image.Image, mode: str, threshold: float = 0.22) -> np.ndarray:
    rgb = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
    if mode == "Green screen":
        # Green dominance rejects green pixels while retaining neutral subject edges.
        score = rgb[:, :, 1] - np.maximum(rgb[:, :, 0], rgb[:, :, 2])
        alpha = 1.0 - np.clip((score - threshold) / max(0.55 - threshold, 1e-3), 0, 1)
    elif mode == "Blue screen":
        score = rgb[:, :, 2] - np.maximum(rgb[:, :, 0], rgb[:, :, 1])
        alpha = 1.0 - np.clip((score - threshold) / max(0.55 - threshold, 1e-3), 0, 1)
    else:
        # Background contrast baseline: estimate foreground as local chroma/saturation.
        chroma = rgb.max(axis=2) - rgb.min(axis=2)
        alpha = np.clip((chroma - threshold) / max(0.65 - threshold, 1e-3), 0, 1)
    # Soft edge prevents harsh cutouts and keeps the result stable for compositing.
    alpha_img = Image.fromarray(np.uint8(alpha * 255), mode="L").filter(ImageFilter.GaussianBlur(1.0))
    return np.asarray(alpha_img, dtype=np.uint8)


def checkerboard(rgb: np.ndarray, alpha: np.ndarray, cell: int = 24) -> Image.Image:
    h, w = alpha.shape
    yy, xx = np.indices((h, w))
    board = np.where(((xx // cell) + (yy // cell)) % 2 == 0, 210, 150).astype(np.float32)
    a = alpha.astype(np.float32) / 255.0
    out = rgb.astype(np.float32) * a[:, :, None] + board[:, :, None] * (1 - a[:, :, None])
    return Image.fromarray(np.uint8(np.clip(out, 0, 255)))


def probe(path: str):
    ffmpeg = __import__("imageio_ffmpeg").get_ffmpeg_exe()
    p = subprocess.run([ffmpeg, "-hide_banner", "-i", path], capture_output=True, text=True)
    m = re.search(r"(\d{2,5})x(\d{2,5})", p.stderr)
    r = re.search(r"(\d+(?:\.\d+)?) fps", p.stderr)
    if not m: raise ValueError("Could not read video dimensions")
    return int(m.group(1)), int(m.group(2)), float(r.group(1)) if r else 24.0


def extract_alpha_video(input_path: str, output_webm: str, preview_mp4: str, mode: str, max_frames: int = 96):
    ffmpeg = __import__("imageio_ffmpeg").get_ffmpeg_exe()
    width, height, fps = probe(input_path)
    scale = min(1.0, 960 / max(width, height)); w, h = max(2, int(width * scale)), max(2, int(height * scale))
    reader = subprocess.Popen([ffmpeg, "-v", "error", "-i", input_path, "-vf", f"scale={w}:{h}", "-f", "image2pipe", "-vcodec", "rawvideo", "-pix_fmt", "rgb24", "-"], stdout=subprocess.PIPE)
    transparent = subprocess.Popen([ffmpeg, "-y", "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgba", "-s", f"{w}x{h}", "-r", str(fps), "-i", "-", "-c:v", "libvpx-vp9", "-pix_fmt", "yuva420p", "-auto-alt-ref", "0", "-b:v", "0", "-crf", "32", output_webm], stdin=subprocess.PIPE)
    preview = subprocess.Popen([ffmpeg, "-y", "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{w}x{h}", "-r", str(fps), "-i", "-", "-c:v", "libx264", "-pix_fmt", "yuv420p", preview_mp4], stdin=subprocess.PIPE)
    size = w * h * 3; count = 0; alpha_sum = 0; first = None; first_alpha = None
    try:
        while count < max_frames:
            raw = reader.stdout.read(size)
            if len(raw) != size: break
            rgb = np.frombuffer(raw, dtype=np.uint8).reshape(h, w, 3).copy()
            alpha = matte_for_frame(Image.fromarray(rgb), mode)
            rgba = np.dstack([rgb, alpha])
            transparent.stdin.write(rgba.tobytes())
            preview.stdin.write(checkerboard(rgb, alpha).tobytes())
            if first is None: first, first_alpha = checkerboard(rgb, alpha), Image.fromarray(alpha, mode="L")
            alpha_sum += float(alpha.mean()) / 255; count += 1
    finally:
        transparent.stdin.close(); preview.stdin.close(); reader.stdout.close()
        transparent.wait(timeout=60); preview.wait(timeout=60); reader.wait(timeout=60)
    if count == 0: raise ValueError("No frames decoded")
    return first, first_alpha, {"project": "SimamAlphaXR", "frames": count, "source_size": [width, height], "processed_size": [w, h], "fps": fps, "mean_alpha_coverage": alpha_sum / count, "matte_mode": mode, "transparent_codec": "VP9 yuva420p", "learned_model_used": False, "claim_boundary": "compositing prototype; not generated alpha-video inference"}


def save_diagnostics(data, path): Path(path).write_text(json.dumps(data, indent=2), encoding="utf-8")