"""Image saving module for saving camera frames in FITS format.

This module provides functionality to save raw YUYV frames as FITS files
with separate Y, U, V channels.
"""

import os
from datetime import datetime
from typing import Optional, Tuple

import cv2
import numpy as np
from astropy.io import fits


def unpack_yuyv(frame_data: bytes, width: int, height: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Unpack YUYV frame data into separate Y, U, V channels.
    
    YUYV format stores pixels as: [Y0, U, Y1, V] for each pair of pixels.
    Y is full resolution, U and V are horizontally subsampled (chroma subsampling 4:2:2).
    
    Args:
        frame_data: Raw YUYV bytes.
        width: Frame width in pixels.
        height: Frame height in pixels.
        
    Returns:
        Tuple of (Y, U, V) arrays where:
        - Y has shape (height, width)
        - U has shape (height, width // 2)
        - V has shape (height, width // 2)
    """
    a = np.frombuffer(frame_data, dtype=np.uint8).reshape(height, width * 2)
    Y = a[:, 0::2].copy()
    U = a[:, 1::4].copy()
    V = a[:, 3::4].copy()
    return Y, U, V


def yuyv_to_rgb(Y: np.ndarray, U: np.ndarray, V: np.ndarray) -> np.ndarray:
    """Convert YUV422 to RGB using BT.601 standard.
    
    Args:
        Y: Luma channel, shape (height, width).
        U: Cb chroma channel, shape (height, width // 2), horizontally subsampled.
        V: Cr chroma channel, shape (height, width // 2), horizontally subsampled.
        
    Returns:
        RGB array with shape (height, width, 3) in uint8 format.
    """
    height, width = Y.shape
    
    # Upsample U and V to full resolution by duplicating each value
    U_full = np.repeat(U, 2, axis=1)
    V_full = np.repeat(V, 2, axis=1)
    
    # Ensure dimensions match after upsampling
    if U_full.shape[1] < width:
        U_full = np.pad(U_full, ((0, 0), (0, width - U_full.shape[1])), mode='edge')
        V_full = np.pad(V_full, ((0, 0), (0, width - V_full.shape[1])), mode='edge')
    elif U_full.shape[1] > width:
        U_full = U_full[:, :width]
        V_full = V_full[:, :width]
    
    # Convert to float and apply BT.601 matrix
    Y_f = Y.astype(np.float32)
    U_f = U_full.astype(np.float32) - 128.0
    V_f = V_full.astype(np.float32) - 128.0
    
    R = Y_f + 1.402 * V_f
    G = Y_f - 0.344136 * U_f - 0.714136 * V_f
    B = Y_f + 1.772 * U_f
    
    RGB = np.stack([R, G, B], axis=-1)
    RGB = np.clip(RGB, 0, 255).astype(np.uint8)
    
    return RGB


def save_fits_image(
    frame_data: bytes,
    width: int,
    height: int,
    hostname: str,
    ts_before: str,
    ts_after: str,
    seq: int,
    checksum: str,
    output_dir: str = ".",
) -> str:
    """Save a YUYV frame as a FITS file with Y, U, V channels.
    
    Filename format: {hostname}_YYmmdd_HHMMSS_fff.fits
    where fff is fraction of seconds (milliseconds).
    
    The FITS file contains:
    - Primary HDU: Y (luma) channel with metadata
    - ImageHDU 1: U (Cb) channel
    - ImageHDU 2: V (Cr) channel  
    
    Header metadata includes timestamps, pixel format, sequence number, and checksum.
    
    Args:
        frame_data: Raw YUYV frame bytes.
        width: Frame width in pixels.
        height: Frame height in pixels.
        hostname: Remote hostname (for filename and header).
        ts_before: Timestamp before capture (ISO8601 format).
        ts_after: Timestamp after capture (ISO8601 format).
        seq: Frame sequence number.
        checksum: MD5 checksum of frame data.
        output_dir: Directory to save the file (default: current directory).
        
    Returns:
        Path to the saved FITS file.
    """
    # Unpack YUYV into separate channels
    Y, U, V = unpack_yuyv(frame_data, width, height)
    
    # Generate filename: hostname_YYmmdd_HHMMSS_fff.fits
    now = datetime.now()
    fraction = int((now.microsecond / 1_000_000) * 1000)
    filename = f"{hostname}_{now.strftime('%y%m%d_%H%M%S')}_{fraction:03d}.fits"
    filepath = os.path.join(output_dir, filename)
    
    # Primary HDU: Y channel with all metadata
    hdu_y = fits.PrimaryHDU(Y)
    hdu_y.header['EXTNAME'] = 'Y'
    hdu_y.header['PIXFMT'] = 'YUYV422'
    hdu_y.header['SEQ'] = seq
    hdu_y.header['HOSTNAME'] = hostname
    hdu_y.header['TS_BEFORE'] = ts_before
    hdu_y.header['TS_AFTER'] = ts_after
    hdu_y.header['CHECKSUM'] = checksum
    hdu_y.header['WIDTH'] = width
    hdu_y.header['HEIGHT'] = height
    
    # Calculate MJD for timestamps (Modified Julian Date)
    try:
        ts_dt = datetime.fromisoformat(ts_after.replace('Z', '+00:00').replace(' ', 'T'))
        year, month, day = ts_dt.year, ts_dt.month, ts_dt.day
        hour = ts_dt.hour + ts_dt.minute/60 + ts_dt.second/3600 + ts_dt.microsecond/3600/1e6
        if month <= 2:
            year -= 1
            month += 12
        A = int(year / 100)
        B = 2 - A + int(A / 4)
        JD = int(365.25 * (year + 4716)) + int(30.6001 * (month + 1)) + day + hour/24 + B - 1524.5
        MJD = JD - 2400000.5
        hdu_y.header['MJD-OBS'] = MJD
    except:
        pass
    
    # ImageHDU for U channel
    hdu_u = fits.ImageHDU(U, name='U')
    
    # ImageHDU for V channel
    hdu_v = fits.ImageHDU(V, name='V')
    
    # Write FITS
    hdul = fits.HDUList([hdu_y, hdu_u, hdu_v])
    hdul.writeto(filepath, overwrite=True)
    
    return filepath


def save_current_frame(
    header: dict,
    frame_data: bytes,
    output_dir: str = ".",
) -> Optional[str]:
    """Convenience function to save a frame from header and data.
    
    Args:
        header: Frame header dictionary with metadata.
        frame_data: Raw frame bytes.
        output_dir: Directory to save the file.
        
    Returns:
        Path to saved file or None if save failed.
    """
    try:
        width = header.get("width", 0)
        height = header.get("height", 0)
        hostname = header.get("hostname", "unknown")
        ts_before = header.get("ts_before", "")
        ts_after = header.get("ts_after", "")
        seq = header.get("seq", 0)
        checksum = header.get("checksum", "")
        fmt = header.get("fmt", "")
        
        if width == 0 or height == 0:
            print("Error: Invalid frame dimensions")
            return None
            
        if fmt != "YUYV":
            print(f"Error: Unknown pixel format '{fmt}'. Only YUYV is supported.")
            return None
            
        filepath = save_fits_image(
            frame_data=frame_data,
            width=width,
            height=height,
            hostname=hostname,
            ts_before=ts_before,
            ts_after=ts_after,
            seq=seq,
            checksum=checksum,
            output_dir=output_dir,
        )
        
        return filepath
    except Exception as e:
        print(f"Error saving frame: {e}")
        return None


def load_yuv_from_fits(filepath: str) -> Optional[Tuple[np.ndarray, np.ndarray, np.ndarray]]:
    """Load Y, U, V channels from a FITS file.
    
    Args:
        filepath: Path to the FITS file.
        
    Returns:
        Tuple of (Y, U, V) arrays or None if loading failed.
    """
    try:
        with fits.open(filepath) as hdul:
            Y = hdul['Y'].data
            U = hdul['U'].data
            V = hdul['V'].data
            return Y, U, V
    except Exception as e:
        print(f"Error loading FITS file: {e}")
        return None


def load_sequence_from_fits(filepath: str) -> Optional[Tuple[np.ndarray, int]]:
    """Load a sequence FITS file (4D RGB cube).
    
    Args:
        filepath: Path to the FITS sequence file.
        
    Returns:
        Tuple of (rgb_cube, nframes) where rgb_cube has shape (N, 3, H, W),
        or None if loading failed.
    """
    try:
        with fits.open(filepath) as hdul:
            data = hdul[0].data
            nframes = hdul[0].header.get('NFRAMES', data.shape[0] if len(data.shape) == 4 else 1)
            return data, nframes
    except Exception as e:
        print(f"Error loading sequence FITS file: {e}")
        return None


def convert_fits_to_image(
    fits_path: str,
    output_path: str,
    format: str = "png",
    quality: int = 95,
    frame: Optional[int] = None,
) -> bool:
    """Convert a FITS file to PNG or JPG.
    
    Supports both individual YUV FITS files and sequence (4D RGB cube) FITS files.
    For sequence files, use frame parameter to select which frame to extract.
    
    Args:
        fits_path: Path to the input FITS file.
        output_path: Path to the output image file.
        format: Output format ('png' or 'jpg').
        quality: JPEG quality (1-100), ignored for PNG.
        frame: Frame index to extract from sequence (0-based). If None and file is 
               a sequence, extracts frame 0.
        
    Returns:
        True if conversion succeeded, False otherwise.
    """
    from PIL import Image
    
    try:
        # First try to load as individual YUV FITS
        yuv = load_yuv_from_fits(fits_path)
        
        if yuv is not None:
            # Individual FITS file with YUV channels
            Y, U, V = yuv
            RGB = yuyv_to_rgb(Y, U, V)
        else:
            # Try to load as sequence FITS (4D RGB cube)
            seq_data = load_sequence_from_fits(fits_path)
            if seq_data is None:
                print(f"Error: Could not load {fits_path} as FITS file")
                return False
            
            rgb_cube, nframes = seq_data
            
            # Select frame
            if frame is None:
                frame_idx = 0
            else:
                frame_idx = frame
            
            if frame_idx < 0 or frame_idx >= nframes:
                print(f"Error: Frame index {frame_idx} out of range (0-{nframes-1})")
                return False
            
            # Extract frame from cube (N, 3, H, W) -> (H, W, 3)
            # Transpose from (3, H, W) to (H, W, 3)
            RGB = rgb_cube[frame_idx].transpose(1, 2, 0)
            print(f"  Extracted frame {frame_idx}/{nframes}")
        
        # Save as image
        img = Image.fromarray(RGB)
        
        if format.lower() == "jpg" or format.lower() == "jpeg":
            img.save(output_path, format="JPEG", quality=quality)
        else:
            img.save(output_path, format="PNG")
        
        return True
    except Exception as e:
        print(f"Error converting {fits_path}: {e}")
        return False


def convert_fits_to_mp4(
    fits_path: str,
    output_path: str,
    fps: int = 30,
    quality: int = 95,
) -> bool:
    """Convert a sequence FITS file to MP4 video.
    
    Args:
        fits_path: Path to the FITS sequence file (4D RGB cube).
        output_path: Path to the output MP4 file.
        fps: Frames per second for the video (default: 30).
        quality: Video quality (affects encoding, default: 95).
        
    Returns:
        True if conversion succeeded, False otherwise.
    """
    try:
        # Load sequence data
        seq_data = load_sequence_from_fits(fits_path)
        if seq_data is None:
            print(f"Error: Could not load {fits_path} as sequence FITS file")
            return False
        
        rgb_cube, nframes = seq_data
        
        if len(rgb_cube.shape) != 4:
            print(f"Error: File is not a sequence (expected 4D data, got {len(rgb_cube.shape)}D)")
            return False
        
        # Get dimensions (N, 3, H, W)
        n_frames, n_channels, height, width = rgb_cube.shape
        
        # OpenCV uses (H, W, 3) format and BGR color space
        # We'll create video writer
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        
        if not out.isOpened():
            print(f"Error: Could not open video writer for {output_path}")
            return False
        
        print(f"  Encoding {n_frames} frames at {fps} fps...")
        
        for i in range(n_frames):
            # Extract frame (3, H, W) -> (H, W, 3)
            frame_rgb = rgb_cube[i].transpose(1, 2, 0)
            
            # Convert RGB to BGR for OpenCV
            frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
            
            # Write frame
            out.write(frame_bgr)
            
            # Progress indicator
            if (i + 1) % 10 == 0 or i == n_frames - 1:
                print(f"    Frame {i + 1}/{n_frames}")
        
        out.release()
        
        print(f"  Saved: {output_path} ({n_frames} frames, {fps} fps)")
        return True
        
    except Exception as e:
        print(f"Error converting {fits_path} to MP4: {e}")
        return False
