#!/usr/bin/env python3
"""
Extract individual frames from grid image, excluding divider ribbons.

NanoBanana generates grids with white or black divider ribbons between cells.
This script detects and excludes those dividers to extract clean frames.

Usage:
    extract_frames_v3.py STORY_NAME [--take TAKE]
"""
import argparse
import sys
import json
from pathlib import Path
from datetime import datetime

from PIL import Image
import numpy as np

# Add lib to path
sys.path.insert(0, str(Path(__file__).parent))

from lib import (
    get_story_path,
    get_take_path,
    get_stage_path,
    ensure_dir,
    update_current_symlink,
    save_metadata,
    log_stage_start,
    log_stage_complete,
    log_stage_error
)


def detect_dividers(img_array, axis=0, threshold=0.95):
    """
    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
    """
    # Sum along opposite axis to get line profiles
    # Average across RGB channels first if image is color
    if len(img_array.shape) == 3:
        img_gray = np.mean(img_array, axis=2)
    else:
        img_gray = img_array

    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 extract_frames_with_dividers(
    grid_path: Path,
    output_dir: Path,
    layout: str = "3x3"
) -> list:
    """
    Extract frames from grid, automatically detecting and excluding dividers.

    Args:
        grid_path: Path to grid image
        output_dir: Directory to save frames
        layout: Grid layout ("3x3" or "4x4")

    Returns:
        List of frame paths
    """
    # Parse layout
    cols, rows = map(int, layout.split('x'))

    print(f"Loading grid: {grid_path}")
    img = Image.open(grid_path)
    width, height = img.size
    print(f"Grid size: {width}x{height}")
    print(f"Layout: {layout} ({rows}x{cols} = {rows*cols} frames)")

    # Convert to numpy for analysis
    img_array = np.array(img)

    # Detect dividers
    print(f"\nDetecting dividers...")
    h_dividers = detect_dividers(img_array, axis=0)  # Horizontal
    v_dividers = detect_dividers(img_array, axis=1)  # Vertical

    print(f"  Horizontal dividers: {len(h_dividers)} found")
    print(f"  Vertical dividers: {len(v_dividers)} found")

    # Calculate frame boundaries excluding dividers
    # Horizontal boundaries (row separations)
    if len(h_dividers) == rows - 1:
        # Expected number of dividers
        h_boundaries = [0]
        for start, end in h_dividers:
            h_boundaries.append(end)
        h_boundaries.append(height)
    else:
        # No dividers or unexpected count - fall back to simple division
        print(f"  Warning: Expected {rows-1} horizontal dividers, found {len(h_dividers)}")
        print(f"  Using simple division instead")
        frame_height = height // rows
        h_boundaries = [i * frame_height for i in range(rows + 1)]

    # Vertical boundaries (column separations)
    if len(v_dividers) == cols - 1:
        # Expected number of dividers
        v_boundaries = [0]
        for start, end in v_dividers:
            v_boundaries.append(end)
        v_boundaries.append(width)
    else:
        # No dividers or unexpected count - fall back to simple division
        print(f"  Warning: Expected {cols-1} vertical dividers, found {len(v_dividers)}")
        print(f"  Using simple division instead")
        frame_width = width // cols
        v_boundaries = [i * frame_width for i in range(cols + 1)]

    print(f"\nFrame boundaries:")
    print(f"  Horizontal: {h_boundaries}")
    print(f"  Vertical: {v_boundaries}")

    # Extract frames
    ensure_dir(output_dir)
    frame_paths = []

    print(f"\nExtracting frames...")
    for row in range(rows):
        for col in range(cols):
            frame_num = row * cols + col + 1

            # Get boundaries for this frame
            left = v_boundaries[col]
            right = v_boundaries[col + 1]
            upper = h_boundaries[row]
            lower = h_boundaries[row + 1]

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

            # Save frame
            frame_path = output_dir / f"frame_{frame_num:02d}.jpg"
            frame.save(frame_path, "JPEG", quality=95)
            frame_paths.append(frame_path)

            print(f"  ✓ Frame {frame_num}: {frame_path.name} ({frame.size[0]}x{frame.size[1]})")

    return frame_paths


def extract_frames(story_name: str, take: str = "current") -> bool:
    """
    Extract frames from grid for V3 workflow.

    Args:
        story_name: Story name
        take: Take number or "current"

    Returns:
        True if successful, False otherwise
    """
    # Get paths
    try:
        story_path = get_story_path(story_name)
        take_path = get_take_path(story_name, take)
    except Exception as e:
        print(f"Error: {e}")
        return False

    # Get take number from path
    if take == "current":
        current_link = story_path / "current"
        if not current_link.is_symlink():
            print(f"Error: No current take set")
            return False
        take_num = current_link.readlink().name
    else:
        take_num = take

    # Get grid path
    grid_stage_path = take_path / "07-grids" / "current"
    if not grid_stage_path.is_symlink():
        print(f"Error: No current grid iteration")
        return False

    grid_iter_path = grid_stage_path.resolve()
    grid_path = grid_iter_path / "grid.jpg"

    if not grid_path.exists():
        print(f"Error: Grid not found at {grid_path}")
        return False

    # Read grid metadata to get layout
    grid_metadata_path = grid_iter_path / "metadata.json"
    if grid_metadata_path.exists():
        with open(grid_metadata_path) as f:
            grid_metadata = json.load(f)
        layout = grid_metadata.get("grid_layout", "3x3")
    else:
        # Default to 3x3
        layout = "3x3"
        print(f"Warning: Grid metadata not found, assuming {layout}")

    # Create frames stage directory with scene subdirectories
    frames_stage_path = take_path / "09-frames"
    ensure_dir(frames_stage_path)

    # Parse layout to get number of scenes
    cols, rows = map(int, layout.split('x'))
    num_scenes = cols * rows

    print(f"Extracting frames from grid:")
    print(f"  Story: {story_name}")
    print(f"  Take: {take_num}")
    print(f"  Grid: {grid_path}")
    print(f"  Layout: {layout} ({num_scenes} scenes)")
    print(f"  Output: {frames_stage_path}")
    print()

    # Load session ID for logging
    session_file = story_path / "session.txt"
    session_id = session_file.read_text().strip() if session_file.exists() else "unknown"

    try:
        # First, extract all frames to a temp directory
        print(f"Loading grid: {grid_path}")
        img = Image.open(grid_path)
        width, height = img.size
        print(f"Grid size: {width}x{height}")
        print(f"Layout: {layout} ({rows}x{cols} = {num_scenes} frames)")

        # Convert to numpy for divider detection
        img_array = np.array(img)

        # Detect dividers
        print(f"\nDetecting dividers...")
        h_dividers = detect_dividers(img_array, axis=0)  # Horizontal
        v_dividers = detect_dividers(img_array, axis=1)  # Vertical

        print(f"  Horizontal dividers: {len(h_dividers)} found")
        print(f"  Vertical dividers: {len(v_dividers)} found")

        # Calculate frame boundaries
        if len(h_dividers) == rows - 1:
            h_boundaries = [0]
            for start, end in h_dividers:
                h_boundaries.append(end)
            h_boundaries.append(height)
        else:
            print(f"  Warning: Expected {rows-1} horizontal dividers, found {len(h_dividers)}")
            print(f"  Using simple division instead")
            frame_height = height // rows
            h_boundaries = [i * frame_height for i in range(rows + 1)]

        if len(v_dividers) == cols - 1:
            v_boundaries = [0]
            for start, end in v_dividers:
                v_boundaries.append(end)
            v_boundaries.append(width)
        else:
            print(f"  Warning: Expected {cols-1} vertical dividers, found {len(v_dividers)}")
            print(f"  Using simple division instead")
            frame_width = width // cols
            v_boundaries = [i * frame_width for i in range(cols + 1)]

        print(f"\nFrame boundaries:")
        print(f"  Horizontal: {h_boundaries}")
        print(f"  Vertical: {v_boundaries}")

        # Extract frames into scene directories
        print(f"\nCreating scene directories and extracting raw frames...")
        for row in range(rows):
            for col in range(cols):
                scene_num = row * cols + col + 1
                scene_dir = frames_stage_path / f"scene_{scene_num:02d}"
                ensure_dir(scene_dir)

                # Get boundaries for this frame
                left = v_boundaries[col]
                right = v_boundaries[col + 1]
                upper = h_boundaries[row]
                lower = h_boundaries[row + 1]

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

                # Trim any remaining white/black edges from the frame
                frame_array = np.array(frame)
                frame_gray = np.mean(frame_array, axis=2) if len(frame_array.shape) == 3 else frame_array
                frame_normalized = frame_gray / 255.0

                # Detect edges to trim (top, bottom, left, right)
                # More aggressive: trim if brightness >0.7 (light) or <0.15 (dark)
                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 to ensure clean edges
                extra_inset = 2
                trim_top += extra_inset
                trim_bottom += extra_inset
                trim_left += extra_inset
                trim_right += extra_inset

                # Apply trim
                fw, fh = frame.size
                frame = frame.crop((trim_left, trim_top, fw - trim_right, fh - trim_bottom))
                print(f"  Scene {scene_num}: Trimmed edges (T:{trim_top} B:{trim_bottom} L:{trim_left} R:{trim_right})")

                # Save as raw_frame.jpg in scene directory
                raw_frame_path = scene_dir / "raw_frame.jpg"
                frame.save(raw_frame_path, "JPEG", quality=95)

                print(f"  ✓ Scene {scene_num}: {raw_frame_path} ({frame.size[0]}x{frame.size[1]})")

        # Create metadata at stage level
        metadata = {
            "stage": "09-frames",
            "take": take_num,
            "timestamp": datetime.now().isoformat(),
            "source_grid": "../07-grids/current/grid.jpg",
            "extraction_method": "PIL_with_divider_detection",
            "layout": layout,
            "scenes_extracted": num_scenes,
            "structure": "scene_XX/raw_frame.jpg",
            "note": "Each scene will have iterations (001/, 002/) for NanoBanana upscaling",
            "dividers_detected": True,
            "h_boundaries": h_boundaries,
            "v_boundaries": v_boundaries
        }

        metadata_path = frames_stage_path / "metadata.json"
        save_metadata(metadata_path, metadata)
        print(f"\n✓ Saved metadata: {metadata_path}")

        print(f"\n{'='*60}")
        print(f"Frame extraction completed!")
        print(f"{'='*60}")
        print(f"Extracted {num_scenes} raw frames from {layout} grid")
        print(f"Output: {frames_stage_path}")
        print(f"\nStructure created:")
        print(f"  09-frames/")
        print(f"    ├── scene_01/raw_frame.jpg")
        print(f"    ├── scene_02/raw_frame.jpg")
        print(f"    └── ...")
        print(f"\nNext step:")
        print(f"  Generate high-quality frames using NanoBanana")
        print(f"  (Create iterations 001/, 002/ in each scene directory)")

        return True

    except Exception as e:
        print(f"\nError extracting frames: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    parser = argparse.ArgumentParser(
        description="Extract frames from grid (V3 workflow)"
    )
    parser.add_argument("story_name", help="Story name")
    parser.add_argument(
        "--take",
        default="current",
        help="Take number (default: current)"
    )

    args = parser.parse_args()

    success = extract_frames(args.story_name, args.take)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
