"""
Frame extraction library for Stage 09.

Extracts individual frames from grid images, automatically detecting
and excluding divider ribbons. Based on V3 extraction logic.
"""
import json
from pathlib import Path
from typing import Tuple, List
from PIL import Image
import numpy as np

from .path_utils_v3 import ensure_dir


def detect_dividers(img_array: np.ndarray, axis: int = 0, threshold: float = 0.95) -> List[Tuple[int, int]]:
    """
    Detect divider lines (white or black ribbons) along an axis.

    Args:
        img_array: Numpy array of image
        axis: 0 for horizontal dividers, 1 for vertical dividers
        threshold: Ratio of pixels that must be same color for a divider

    Returns:
        List of (start, end) tuples for each divider position
    """
    # Convert to grayscale if color image
    if len(img_array.shape) == 3:
        img_gray = np.mean(img_array, axis=2)
    else:
        img_gray = img_array

    # Calculate line profile (average along opposite axis)
    if axis == 0:
        # Horizontal dividers - average along width
        profile = np.mean(img_gray, axis=1)
    else:
        # Vertical dividers - average along height
        profile = np.mean(img_gray, axis=0)

    # Normalize to 0-1
    profile_normalized = profile / 255.0

    # Detect lines that are very white (>0.9) or very black (<0.1)
    is_divider = (profile_normalized > 0.9) | (profile_normalized < 0.1)

    # Find continuous divider regions
    dividers = []
    in_divider = False
    start = None

    for i, is_div in enumerate(is_divider):
        if is_div and not in_divider:
            # Start of divider
            start = i
            in_divider = True
        elif not is_div and in_divider:
            # End of divider
            end = i
            dividers.append((start, end))
            in_divider = False

    # Handle divider at edge
    if in_divider:
        dividers.append((start, len(is_divider)))

    return dividers


def calculate_boundaries(
    dividers: List[Tuple[int, int]],
    total_size: int,
    grid_count: int
) -> List[int]:
    """
    Calculate frame boundaries from detected dividers.

    Args:
        dividers: List of (start, end) tuples for dividers
        total_size: Total size (width or height) of grid image
        grid_count: Number of rows or columns

    Returns:
        List of boundary positions (length = grid_count + 1)
    """
    expected_divider_count = grid_count - 1

    if len(dividers) == expected_divider_count:
        # Expected number of dividers detected
        boundaries = [0]
        for start, end in dividers:
            boundaries.append(end)
        boundaries.append(total_size)
    else:
        # Fallback to simple proportional division
        frame_size = total_size // grid_count
        boundaries = [i * frame_size for i in range(grid_count + 1)]

    return boundaries


def trim_edges(frame: Image.Image) -> Image.Image:
    """
    Trim white/black edges from frame (divider artifacts).

    Uses aggressive trimming: removes pixels with brightness >0.7 (light)
    or <0.15 (dark), checking up to 50 pixels on each edge.

    Args:
        frame: PIL Image

    Returns:
        Trimmed PIL Image
    """
    frame_array = np.array(frame)

    # Convert to grayscale
    if len(frame_array.shape) == 3:
        frame_gray = np.mean(frame_array, axis=2)
    else:
        frame_gray = frame_array

    frame_normalized = frame_gray / 255.0

    # Initialize trim amounts
    trim_top = 0
    trim_bottom = 0
    trim_left = 0
    trim_right = 0

    # Check top edge (up to 50 pixels)
    for i in range(min(50, frame_array.shape[0])):
        row_mean = np.mean(frame_normalized[i, :])
        if row_mean > 0.7 or row_mean < 0.15:
            trim_top = i + 1
        else:
            break

    # Check bottom edge
    for i in range(min(50, frame_array.shape[0])):
        row_idx = frame_array.shape[0] - 1 - i
        row_mean = np.mean(frame_normalized[row_idx, :])
        if row_mean > 0.7 or row_mean < 0.15:
            trim_bottom = i + 1
        else:
            break

    # Check left edge
    for i in range(min(50, frame_array.shape[1])):
        col_mean = np.mean(frame_normalized[:, i])
        if col_mean > 0.7 or col_mean < 0.15:
            trim_left = i + 1
        else:
            break

    # Check right edge
    for i in range(min(50, frame_array.shape[1])):
        col_idx = frame_array.shape[1] - 1 - i
        col_mean = np.mean(frame_normalized[:, col_idx])
        if col_mean > 0.7 or col_mean < 0.15:
            trim_right = i + 1
        else:
            break

    # Add 2 extra pixels for clean edges
    extra_inset = 2
    trim_top += extra_inset
    trim_bottom += extra_inset
    trim_left += extra_inset
    trim_right += extra_inset

    # Apply trim (with safety check)
    fw, fh = frame.size
    left = trim_left
    top = trim_top
    right = fw - trim_right
    bottom = fh - trim_bottom
    # If trim would produce invalid crop, return untrimmed
    if right <= left or bottom <= top:
        return frame
    return frame.crop((left, top, right, bottom))


def extract_all_frames_from_grids(take_path: Path, skip_grids: list = None) -> dict:
    """
    Extract frames from ALL grid iterations in 07-grids (or only new ones).

    Args:
        take_path: Path to take directory
        skip_grids: List of grid iteration numbers to skip (for incremental extraction)

    Returns:
        {
            "grid_iterations": ["001", "002", "057"],
            "scene_count": 9,
            "layout": "3x3",
            "h_boundaries": [...],
            "v_boundaries": [...],
            "skipped_count": 0,
            "extracted_count": 57
        }
    """
    grids_dir = take_path / "07-grids"
    frames_dir = take_path / "08-frames"
    skip_grids = skip_grids or []

    if not grids_dir.exists():
        raise ValueError(f"Grids directory not found: {grids_dir}")

    # Scan for all grid iterations
    all_grid_iterations = []
    for item in grids_dir.iterdir():
        if item.is_dir() and item.name.isdigit() and len(item.name) == 3:
            grid_file = item / "grid.jpg"
            if not grid_file.exists():
                grid_file = item / "grid.png"
            if grid_file.exists():
                all_grid_iterations.append(item.name)

    all_grid_iterations.sort()

    if not all_grid_iterations:
        raise ValueError("No grid iterations found in 07-grids")

    # Filter out already extracted grids
    grid_iterations = [g for g in all_grid_iterations if g not in skip_grids]

    if not grid_iterations and skip_grids:
        # All grids already extracted
        print(f"All {len(all_grid_iterations)} grids already extracted. Nothing to do.")
        # Still need to return structure, so use first grid for layout
        first_meta = grids_dir / all_grid_iterations[0] / "metadata.json"
        layout = json.load(open(first_meta)).get("grid_layout", "3x3")
        cols, rows = map(int, layout.split('x'))
        return {
            "grid_iterations": all_grid_iterations,
            "scene_count": rows * cols,
            "layout": layout,
            "h_boundaries": [],
            "v_boundaries": [],
            "skipped_count": len(skip_grids),
            "extracted_count": 0
        }

    # Read layout from first grid's metadata
    first_grid_meta = grids_dir / grid_iterations[0] / "metadata.json"
    with open(first_grid_meta) as f:
        meta = json.load(f)
    layout = meta.get("grid_layout", "3x3")

    # Parse layout
    cols_str, rows_str = layout.split('x')
    cols = int(cols_str)
    rows = int(rows_str)
    scene_count = rows * cols

    if skip_grids:
        print(f"Total grids: {len(all_grid_iterations)}")
        print(f"Already extracted: {len(skip_grids)}")
        print(f"New grids to extract: {len(grid_iterations)}")
    else:
        print(f"Extracting frames from {len(grid_iterations)} grid iterations ({layout} layout)")

    print(f"Frames to extract: {len(grid_iterations)} grids × {scene_count} scenes = {len(grid_iterations) * scene_count}")

    # Create frames directory
    ensure_dir(frames_dir)

    # Create scene directories
    for scene_num in range(1, scene_count + 1):
        scene_dir = frames_dir / f"{scene_num:02d}-scene"
        ensure_dir(scene_dir)

    # Variables to store boundaries from first grid (for metadata)
    h_boundaries_final = None
    v_boundaries_final = None

    # Extract from each grid iteration
    for grid_idx, grid_iter in enumerate(grid_iterations, 1):
        print(f"\n[{grid_idx}/{len(grid_iterations)}] Processing grid {grid_iter}...")

        # Find grid image
        grid_path = grids_dir / grid_iter / "grid.jpg"
        if not grid_path.exists():
            grid_path = grids_dir / grid_iter / "grid.png"

        # Load image
        img = Image.open(grid_path)
        width, height = img.size
        img_array = np.array(img)

        # Detect dividers
        h_dividers = detect_dividers(img_array, axis=0)
        v_dividers = detect_dividers(img_array, axis=1)

        # Calculate boundaries
        h_boundaries = calculate_boundaries(h_dividers, height, rows)
        v_boundaries = calculate_boundaries(v_dividers, width, cols)

        # Store from first grid for metadata
        if grid_idx == 1:
            h_boundaries_final = h_boundaries
            v_boundaries_final = v_boundaries

        # Extract each scene frame
        for row in range(rows):
            for col in range(cols):
                scene_num = row * cols + col + 1
                scene_dir = frames_dir / f"{scene_num:02d}-scene"

                # Crop frame
                left = v_boundaries[col]
                right = v_boundaries[col + 1]
                upper = h_boundaries[row]
                lower = h_boundaries[row + 1]

                frame = img.crop((left, upper, right, lower))

                # Trim edges to remove divider artifacts
                frame = trim_edges(frame)

                # Skip frames that are too small (bad grid layout)
                if frame.size[0] < 10 or frame.size[1] < 10:
                    print(f"  ⚠ Skipping scene {scene_num} from grid {grid_iter} (frame too small: {frame.size})")
                    continue

                # Save as {grid_iter}.jpg
                frame_path = scene_dir / f"{grid_iter}.jpg"
                frame.save(frame_path, "JPEG", quality=95)

        print(f"  ✓ Extracted {scene_count} frames from grid {grid_iter}")

    return {
        "grid_iterations": all_grid_iterations,  # Return ALL grids (old + new)
        "scene_count": scene_count,
        "layout": layout,
        "h_boundaries": h_boundaries_final,
        "v_boundaries": v_boundaries_final,
        "skipped_count": len(skip_grids),
        "extracted_count": len(grid_iterations)
    }
