"""FITS image viewer for viewing saved camera frames.

This module provides functionality to view FITS files created by the image_saving
module, displaying Y, U, V channels and RGB preview.
"""

import os
import sys
import select
import termios
import tty
from pathlib import Path
from typing import Optional, List, Tuple

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


def get_key() -> str:
    """Read a single keypress from terminal.
    
    Returns:
        Key name ('up', 'down', 'enter', 'q', etc.) or the character.
    """
    # Save terminal settings
    old_settings = termios.tcgetattr(sys.stdin)
    
    try:
        # Set terminal to raw mode
        tty.setcbreak(sys.stdin.fileno())
        
        # Wait for input
        if select.select([sys.stdin], [], [], 0.1)[0]:
            char = sys.stdin.read(1)
            
            # Check for escape sequences (arrow keys)
            if char == '\x1b':  # ESC
                # Read next two characters for arrow keys
                next_chars = sys.stdin.read(2)
                if next_chars == '[A':
                    return 'up'
                elif next_chars == '[B':
                    return 'down'
                elif next_chars == '':
                    return 'esc'
                return f'esc+{next_chars}'
            elif char == '\n' or char == '\r':
                return 'enter'
            elif char == '\x03':  # Ctrl+C
                return 'ctrl_c'
            elif char == 'q':
                return 'q'
            elif char == 'j':
                return 'down'
            elif char == 'k':
                return 'up'
            else:
                return char
        
        return ''
    finally:
        # Restore terminal settings
        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings)


def interactive_file_selection(file_list: List[str], directory: str) -> Optional[str]:
    """Display interactive menu for file selection using arrow keys.
    
    Args:
        file_list: List of file paths to select from.
        directory: Directory where files are located.
        
    Returns:
        Selected file path or None if user quit.
    """
    # Load file info for all files
    file_info_list = []
    for filepath in file_list:
        hdul = load_fits_file(filepath)
        if hdul:
            info = get_fits_info(hdul)
            file_info_list.append({
                'path': filepath,
                'filename': info['filename'],
                'hostname': info['hostname'],
                'seq': info['seq'],
                'time': info['ts_after']
            })
            hdul.close()
    
    if not file_info_list:
        print("No valid FITS files to display")
        return None
    
    selected_idx = 0
    
    def print_menu():
        """Print the menu with current selection highlighted."""
        # Clear screen (ANSI escape sequence)
        sys.stdout.write('\033[2J\033[H')
        sys.stdout.flush()
        
        print(f"\nFound {len(file_info_list)} FITS file(s) in {directory}:\n")
        print("Use ↑/↓ arrows (or k/j) to navigate, Enter to select, q to quit\n")
        
        for i, info in enumerate(file_info_list):
            cursor = "> " if i == selected_idx else "  "
            highlight_start = "\033[7m" if i == selected_idx else ""  # Reverse video
            highlight_end = "\033[0m" if i == selected_idx else ""     # Reset
            print(f"{cursor}{highlight_start}{i+1}. {info['filename']}  Seq:{info['seq']}{highlight_end}")
    
    # Initial display
    print_menu()
    
    while True:
        key = get_key()
        
        if key == 'up':
            selected_idx = max(0, selected_idx - 1)
            print_menu()
        elif key == 'down':
            selected_idx = min(len(file_info_list) - 1, selected_idx + 1)
            print_menu()
        elif key == 'enter':
            # Clear screen before returning
            sys.stdout.write('\033[2J\033[H')
            sys.stdout.flush()
            return file_info_list[selected_idx]['path']
        elif key == 'q' or key == 'esc' or key == 'ctrl_c':
            # Clear screen before quitting
            sys.stdout.write('\033[2J\033[H')
            sys.stdout.flush()
            return None


def load_fits_file(filepath: str) -> Optional[fits.HDUList]:
    """Load a FITS file and return the HDUList.
    
    Args:
        filepath: Path to the FITS file.
        
    Returns:
        HDUList if successful, None otherwise.
    """
    try:
        hdul = fits.open(filepath)
        return hdul
    except Exception as e:
        print(f"Error loading FITS file: {e}")
        return None


def get_fits_info(hdul: fits.HDUList) -> dict:
    """Extract information from FITS file headers.
    
    Args:
        hdul: Opened FITS HDUList.
        
    Returns:
        Dictionary with file information.
    """
    info = {}
    primary = hdul[0]
    
    info['filename'] = os.path.basename(hdul.filename())
    info['shape'] = primary.data.shape
    info['dtype'] = primary.data.dtype
    info['ext_count'] = len(hdul)
    
    # Extract header metadata
    header = primary.header
    info['hostname'] = header.get('HOSTNAME', 'unknown')
    info['seq'] = header.get('SEQ', 0)
    info['pixfmt'] = header.get('PIXFMT', 'unknown')
    info['ts_before'] = header.get('TS_BEFORE', 'N/A')
    info['ts_after'] = header.get('TS_AFTER', 'N/A')
    info['checksum'] = header.get('CHECKSUM', 'N/A')
    info['mjd_obs'] = header.get('MJD-OBS', None)
    
    # List available extensions
    info['extensions'] = []
    for i, hdu in enumerate(hdul):
        extname = hdu.header.get('EXTNAME', f'HDU_{i}')
        info['extensions'].append((i, extname, hdu.data.shape if hdu.data is not None else None))
    
    return info


def normalize_for_display(data: np.ndarray) -> np.ndarray:
    """Normalize array to 0-255 range for display.
    
    Args:
        data: Input array.
        
    Returns:
        Normalized uint8 array.
    """
    if data.dtype == np.uint8:
        return data
    
    # Normalize to 0-255
    data_min = data.min()
    data_max = data.max()
    
    if data_max > data_min:
        normalized = ((data - data_min) / (data_max - data_min) * 255).astype(np.uint8)
    else:
        normalized = np.zeros_like(data, dtype=np.uint8)
    
    return normalized


def get_jpg_from_fits(hdul: fits.HDUList) -> Optional[np.ndarray]:
    """Extract and decompress JPG data from FITS file.
    
    Args:
        hdul: Opened FITS HDUList.
        
    Returns:
        BGR image array or None if failed.
    """
    try:
        # First try to find JPG in ImageHDU (new compressed format)
        for hdu in hdul:
            if hdu.header.get('EXTNAME') == 'JPG' and hdu.data is not None:
                # JPG is stored as 1D uint8 array
                jpg_bytes = bytes(hdu.data)
                # Decompress JPG
                img_array = np.frombuffer(jpg_bytes, dtype=np.uint8)
                img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
                if img is not None:
                    return img
        
        # Fallback: try to find raw RGB (old format)
        for hdu in hdul:
            if hdu.header.get('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 to (H, W, C) and convert to BGR
                    rgb_hwc = np.transpose(rgb_data, (1, 2, 0))
                    bgr_data = cv2.cvtColor(rgb_hwc, cv2.COLOR_RGB2BGR)
                    return bgr_data
        
        print("Error: No JPG or RGB data found in FITS file")
        return None
        
    except Exception as e:
        print(f"Error extracting image from FITS: {e}")
        return None


def save_fits_as_image(hdul: fits.HDUList, filepath: str, img_format: str) -> Optional[str]:
    """Save FITS image as JPG or PNG.
    
    Args:
        hdul: Opened FITS HDUList.
        filepath: Original FITS file path (for filename base).
        img_format: 'jpg' or 'png'.
        
    Returns:
        Path to saved image file or None if failed.
    """
    try:
        # Get image data (decompress JPG if needed)
        img_bgr = get_jpg_from_fits(hdul)
        if img_bgr is None:
            return None
        
        # Check if FITS has JPG extension and we're saving JPG
        # If so, we can just extract the bytes directly without re-encoding
        if img_format.lower() == 'jpg':
            for hdu in hdul:
                if hdu.header.get('EXTNAME') == 'JPG' and hdu.data is not None:
                    # Direct copy of JPG bytes - no re-encoding!
                    jpg_bytes = bytes(hdu.data)
                    base_name = os.path.splitext(os.path.basename(filepath))[0]
                    output_filename = f"{base_name}.jpg"
                    with open(output_filename, 'wb') as f:
                        f.write(jpg_bytes)
                    return output_filename
        
        # Otherwise encode from BGR array
        base_name = os.path.splitext(os.path.basename(filepath))[0]
        ext = img_format.lower()
        output_filename = f"{base_name}.{ext}"
        
        if ext == 'jpg':
            cv2.imwrite(output_filename, img_bgr, [cv2.IMWRITE_JPEG_QUALITY, 95])
        else:  # png
            cv2.imwrite(output_filename, img_bgr)
        
        return output_filename
        
    except Exception as e:
        print(f"Error saving image: {e}")
        return None


def display_fits_channels(hdul: fits.HDUList, scale: float = 1.0) -> None:
    """Display image from FITS file.
    
    Args:
        hdul: Opened FITS HDUList.
        scale: Display scaling factor (default: 1.0).
    """
    window_names = []
    
    try:
        # Get image from JPG or RGB
        img_bgr = get_jpg_from_fits(hdul)
        
        if img_bgr is None:
            return []
        
        # Apply scaling
        if scale != 1.0:
            new_width = int(img_bgr.shape[1] * scale)
            new_height = int(img_bgr.shape[0] * scale)
            img_bgr = cv2.resize(img_bgr, (new_width, new_height), interpolation=cv2.INTER_LINEAR)
        
        # Create window
        window_name = f"Image - {os.path.basename(hdul.filename())}"
        cv2.imshow(window_name, img_bgr)
        window_names.append(window_name)
        
        # Add info window
        info = get_fits_info(hdul)
        info_text = [
            f"File: {info['filename']}",
            f"Hostname: {info['hostname']}",
            f"Sequence: {info['seq']}",
            f"Format: {info['pixfmt']}",
            f"Shape: {info['shape']}",
            f"T1: {info['ts_before']}",
            f"T2: {info['ts_after']}",
            f"Extensions: {info['ext_count']}",
        ]
        
        # Create info image
        line_height = 25
        info_img = np.ones((len(info_text) * line_height + 20, 400, 3), dtype=np.uint8) * 240
        for j, line in enumerate(info_text):
            cv2.putText(info_img, line, (10, 30 + j * line_height), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1, cv2.LINE_AA)
        
        cv2.imshow("FITS Info", info_img)
        window_names.append("FITS Info")
        
        return window_names
        
    except Exception as e:
        print(f"Error displaying channels: {e}")
        return []


def close_windows(window_names: List[str]) -> None:
    """Close all OpenCV windows.
    
    Args:
        window_names: List of window names to close.
    """
    for name in window_names:
        cv2.destroyWindow(name)


def view_fits_file(filepath: str, scale: float = 1.0) -> bool:
    """View a single FITS file.
    
    Args:
        filepath: Path to FITS file.
        scale: Display scaling factor.
        
    Returns:
        True if user wants to continue, False to quit.
    """
    hdul = load_fits_file(filepath)
    if hdul is None:
        return True  # Continue to next file
    
    try:
        info = get_fits_info(hdul)
        print(f"\nViewing: {info['filename']}")
        print(f"  Hostname: {info['hostname']}, Seq: {info['seq']}")
        print(f"  Shape: {info['shape']}, Format: {info['pixfmt']}")
        print(f"  Time: {info['ts_after']}")
        print(f"  Extensions: {[ext[1] for ext in info['extensions']]}")
        
        window_names = display_fits_channels(hdul, scale)
        
        print("  Keys: q=menu, j=JPG, p=PNG, 1-4=scale")
        
        current_scale = scale
        while True:
            key = cv2.waitKey(100) & 0xFF
            
            if key == ord('q') or key == 27:  # q or ESC - return to menu
                close_windows(window_names)
                return True  # Return to menu
            elif key == ord('j'):
                # Save as JPG
                output = save_fits_as_image(hdul, filepath, 'jpg')
                if output:
                    print(f"  Saved as JPG: {output}")
            elif key == ord('p'):
                # Save as PNG
                output = save_fits_as_image(hdul, filepath, 'png')
                if output:
                    print(f"  Saved as PNG: {output}")
            elif key == ord('1'):
                current_scale = 1.0
                close_windows(window_names)
                window_names = display_fits_channels(hdul, current_scale)
            elif key == ord('2'):
                current_scale = 2.0
                close_windows(window_names)
                window_names = display_fits_channels(hdul, current_scale)
            elif key == ord('3'):
                current_scale = 0.5
                close_windows(window_names)
                window_names = display_fits_channels(hdul, current_scale)
            elif key == ord('4'):
                current_scale = 0.25
                close_windows(window_names)
                window_names = display_fits_channels(hdul, current_scale)
                
    finally:
        hdul.close()
        cv2.destroyAllWindows()
    
    return True


def find_fits_files(directory: str = ".") -> List[str]:
    """Find all FITS files in directory.
    
    Args:
        directory: Directory to search (default: current directory).
        
    Returns:
        List of FITS file paths sorted by modification time.
    """
    fits_extensions = ['.fits', '.fit', '.fts']
    files = []
    
    for ext in fits_extensions:
        files.extend(Path(directory).glob(f'*{ext}'))
    
    # Sort by modification time (oldest first)
    files.sort(key=lambda x: x.stat().st_mtime)
    
    return [str(f) for f in files]


@click.command()
@click.argument('files', nargs=-1, required=False)
@click.option(
    '--directory', '-d',
    default='.',
    help='Directory to search for FITS files (default: current directory)',
)
@click.option(
    '--scale', '-s',
    default=1.0,
    type=float,
    help='Display scaling factor (default: 1.0)',
)
@click.option(
    '--list-only', '-l',
    is_flag=True,
    help='List all FITS files without viewing',
)
def view_fits(files: Tuple[str], directory: str, scale: float, list_only: bool) -> None:
    """View FITS image files saved by zmqcam.
    
    If no files are specified, searches the current directory for .fits files.
    
    Examples:
        zmqcam-view fits_file.fits
        zmqcam-view -d /path/to/images
        zmqcam-view -l                    # List all FITS files
        zmqcam-view *.fits --scale 2.0    # View with 2x scaling
    """
    # Find files to view
    file_list = list(files)
    
    if not file_list:
        file_list = find_fits_files(directory)
        if not file_list:
            print(f"No FITS files found in {directory}")
            sys.exit(1)
        
        # Interactive file selection with arrow keys
        selected_file = interactive_file_selection(file_list, directory)
        if selected_file is None:
            sys.exit(0)
        file_list = [selected_file]
    
    # Just list files if requested
    if list_only:
        for i, filepath in enumerate(file_list, 1):
            hdul = load_fits_file(filepath)
            if hdul:
                info = get_fits_info(hdul)
                print(f"{i}. {info['filename']}")
                print(f"   Hostname: {info['hostname']}, Seq: {info['seq']}, Time: {info['ts_after']}")
                print(f"   Shape: {info['shape']}, Extensions: {info['ext_count']}")
                hdul.close()
        return
    
    # View files interactively
    current_index = 0
    
    while 0 <= current_index < len(file_list):
        filepath = file_list[current_index]
        
        result = view_fits_file(filepath, scale)
        
        if result == False:  # Quit
            break
        elif result == 'previous':  # Previous file
            current_index = max(0, current_index - 1)
        else:  # Next file
            current_index += 1
            if current_index >= len(file_list):
                print("\nReached end of file list.")
                break
    
    cv2.destroyAllWindows()
    print("Viewer closed.")


if __name__ == '__main__':
    view_fits()
