"""Napari-based FITS file viewer for viewing stacked camera frames.

This module provides functionality to combine multiple FITS files into a
multi-page TIFF stack and view them interactively using napari.
"""

import os
import sys
import tempfile
from pathlib import Path
from typing import List, Optional, Tuple
import glob

import click
import numpy as np
from astropy.io import fits
import tifffile
from natsort import natsorted


def find_fits_files(patterns: List[str]) -> List[str]:
    """Find FITS files matching the given patterns.
    
    Args:
        patterns: List of file patterns (e.g., ['*.fits', 'image_*.fit']).
        
    Returns:
        List of matching file paths, naturally sorted.
    """
    files = []
    fits_extensions = ['.fits', '.fit', '.fts']
    
    for pattern in patterns:
        # Check if pattern is a specific file
        path = Path(pattern)
        if path.is_file():
            files.append(str(path))
        else:
            # Use glob pattern
            matches = glob.glob(pattern)
            files.extend(matches)
    
    # Filter to only include FITS files
    files = [f for f in files if any(f.lower().endswith(ext) for ext in fits_extensions)]
    
    # Remove duplicates and sort naturally
    files = list(set(files))
    files = natsorted(files)
    
    return files


def load_fits_data(filepath: str) -> Optional[np.ndarray]:
    """Load image data from a FITS file.
    
    Args:
        filepath: Path to the FITS file.
        
    Returns:
        Numpy array containing the image data, or None if failed.
    """
    try:
        with fits.open(filepath) as hdul:
            # Try to get RGB or JPG data first
            for hdu in hdul:
                extname = hdu.header.get('EXTNAME', '')
                if extname == 'RGB' and hdu.data is not None:
                    rgb_data = hdu.data
                    if len(rgb_data.shape) == 3 and rgb_data.shape[0] == 3:
                        # Transpose from (C, H, W) to (H, W, C)
                        return np.transpose(rgb_data, (1, 2, 0))
                    return rgb_data
            
            # Fall back to primary data
            primary_data = hdul[0].data
            if primary_data is not None:
                return primary_data
                
        return None
    except Exception as e:
        print(f"Error loading {filepath}: {e}")
        return None


def create_tiff_stack(
    files: List[str],
    output_path: str,
    dtype: Optional[str] = None
) -> bool:
    """Combine FITS files into a multi-page TIFF stack.
    
    Args:
        files: List of FITS file paths.
        output_path: Path for the output TIFF file.
        dtype: Optional dtype to cast data to (e.g., 'float32').
        
    Returns:
        True if successful, False otherwise.
    """
    if not files:
        print("No FITS files to process")
        return False
    
    frames = []
    print(f"Processing {len(files)} FITS files...")
    
    for i, filepath in enumerate(files):
        print(f"  [{i+1}/{len(files)}] {os.path.basename(filepath)}")
        data = load_fits_data(filepath)
        if data is not None:
            frames.append(data)
    
    if not frames:
        print("No valid frames loaded")
        return False
    
    try:
        # Stack along first axis (time/stack dimension)
        stack = np.stack(frames, axis=0)
        
        # Optionally cast to smaller dtype
        if dtype:
            stack = stack.astype(dtype)
        
        # Write multi-page TIFF
        tifffile.imwrite(output_path, stack)
        print(f"Created TIFF stack: {output_path}")
        print(f"  Shape: {stack.shape}")
        print(f"  Dtype: {stack.dtype}")
        return True
        
    except Exception as e:
        print(f"Error creating TIFF stack: {e}")
        return False


def view_with_napari(tiff_path: str) -> None:
    """Open a TIFF stack in napari viewer.
    
    Args:
        tiff_path: Path to the TIFF file to view.
    """
    # Fix Qt platform plugin conflict with opencv-python
    # Remove cv2's Qt plugins from path to avoid conflicts with napari's Qt
    if 'QT_QPA_PLATFORM_PLUGIN_PATH' in os.environ:
        del os.environ['QT_QPA_PLATFORM_PLUGIN_PATH']
    
    # Set platform to wayland if XDG_SESSION_TYPE is wayland, otherwise xcb
    if os.environ.get('XDG_SESSION_TYPE') == 'wayland':
        os.environ['QT_QPA_PLATFORM'] = 'wayland'
    else:
        os.environ['QT_QPA_PLATFORM'] = 'xcb'
    
    try:
        import napari
        
        print(f"Loading TIFF stack into napari...")
        data = tifffile.imread(tiff_path)
        
        viewer = napari.Viewer()
        viewer.add_image(data, name="FITS Stack")
        
        print("Starting napari viewer (close window to exit)...")
        napari.run()
        
    except ImportError:
        print("Error: napari is not installed. Install with: pip install napari")
        sys.exit(1)
    except Exception as e:
        print(f"Error viewing with napari: {e}")
        sys.exit(1)


def generate_temp_filename(files: List[str]) -> str:
    """Generate a unique temporary filename based on first FITS file, count, and username.
    
    Args:
        files: List of FITS file paths.
        
    Returns:
        Unique filename string.
    """
    import getpass
    
    if not files:
        return "napari_fits_stack.tif"
    
    # Get first filename without extension
    first_file = os.path.basename(files[0])
    first_name = os.path.splitext(first_file)[0]
    
    # Get number of files
    n_files = len(files)
    
    # Get username
    try:
        username = getpass.getuser()
    except:
        username = "unknown"
    
    return f"{first_name}_{n_files}_{username}.tif"


def process_and_view(
    patterns: List[str],
    keep_temp: bool = False,
    temp_dir: Optional[str] = None
) -> None:
    """Process FITS files and view them with napari.
    
    Args:
        patterns: List of file patterns to match.
        keep_temp: If True, keep the temporary TIFF file.
        temp_dir: Optional directory for temporary file (default: /tmp).
    """
    # Find matching files
    files = find_fits_files(patterns)
    
    if not files:
        print(f"No FITS files found matching patterns: {patterns}")
        sys.exit(1)
    
    print(f"Found {len(files)} FITS file(s)")
    
    # Create temporary TIFF file with unique name
    temp_dir = temp_dir or "/tmp"
    temp_filename = generate_temp_filename(files)
    temp_path = os.path.join(temp_dir, temp_filename)
    
    try:
        # Create the TIFF stack
        if not create_tiff_stack(files, temp_path):
            sys.exit(1)
        
        # View with napari
        view_with_napari(temp_path)
        
    finally:
        # Clean up temporary file
        if not keep_temp and os.path.exists(temp_path):
            os.remove(temp_path)
            print(f"Cleaned up temporary file: {temp_path}")


@click.command()
@click.argument('patterns', nargs=-1, required=True)
@click.option(
    '--keep-temp', '-k',
    is_flag=True,
    help='Keep the temporary TIFF file after viewing',
)
@click.option(
    '--temp-dir', '-t',
    default='/tmp',
    help='Directory for temporary TIFF file (default: /tmp)',
)
def napari_cli(patterns: Tuple[str], keep_temp: bool, temp_dir: str) -> None:
    """View FITS files as a stack in napari.
    
    Combines multiple FITS files into a temporary multi-page TIFF stack
    and opens it in napari for interactive viewing.
    
    Examples:
        zmqcam napari *.fits
        zmqcam napari image_*.fits
        zmqcam napari --keep-temp *.fits
        zmqcam napari --temp-dir /var/tmp *.fits
    """
    process_and_view(list(patterns), keep_temp=keep_temp, temp_dir=temp_dir)


if __name__ == '__main__':
    napari_cli()
