//! CLI functionality shared between the binary and Python interface

use crate::{CompressionSettings, PrecisionMode, Reader, Writer};
use anyhow::Result;
use clap::{Parser, Subcommand};

use colored::*;
use std::path::PathBuf;

/// Print a success message
fn print_success(msg: &str) {
    println!("{} {}", "✔️".green(), msg.green());
}

/// Print an error message
fn print_error(msg: &str) {
    eprintln!("{} {}", "✖️".red(), msg.red());
}

/// Print a warning message
#[allow(dead_code)]
fn print_warning(msg: &str) {
    println!("{} {}", "⚠️".yellow(), msg.yellow());
}

/// Print an info/progress message
fn print_info(msg: &str) {
    println!("{} {}", "➜".cyan(), msg.cyan());
}

/// TACO trajectory file utility
#[derive(Parser)]
#[command(name = "taco")]
#[command(about = "TACO trajectory file utility")]
#[command(version)]
pub struct Cli {
    #[command(subcommand)]
    pub command: Commands,
}

#[derive(Subcommand)]
pub enum Commands {
    /// Display information about a TACO file
    Info {
        /// Path to the TACO file
        input: PathBuf,
    },
    /// Copy a subset of frames to a new file
    Extract {
        /// Input TACO file
        input: PathBuf,
        /// Output TACO file
        output: PathBuf,
        /// Start frame (0-based, inclusive)
        #[arg(short, long)]
        start: Option<u64>,
        /// End frame (0-based, exclusive)
        #[arg(short, long)]
        end: Option<u64>,
        /// Specific frame indices to extract (comma-separated)
        #[arg(short, long, value_delimiter = ',')]
        frames: Option<Vec<u64>>,
        /// Skip every N frames (extract every Nth frame)
        #[arg(long)]
        step: Option<u64>,
    },
    /// Convert between different trajectory formats
    Convert {
        /// Input trajectory file
        input: PathBuf,
        /// Output TACO file
        output: PathBuf,
        /// Compression level (1-22, higher = better compression)
        #[arg(short, long, default_value = "3")]
        compression: i32,
        /// Use full precision (32-bit float)
        #[arg(long)]
        full_precision: bool,
        /// Use lossless compression (alias for full-precision)
        #[arg(long)]
        lossless: bool,
        /// Full frame interval (store complete frame every N frames)
        #[arg(long, default_value = "10")]
        full_frame_interval: u32,
    },
    /// Validate the integrity of a TACO file
    Check {
        /// Path to the TACO file
        input: PathBuf,
        /// Verbose output
        #[arg(short, long)]
        verbose: bool,
    },
    /// Show statistics about a TACO file
    Stats {
        /// Path to the TACO file
        input: PathBuf,
        /// Show frame-by-frame statistics
        #[arg(short, long)]
        detailed: bool,
    },
}

/// Execute a CLI command
pub fn execute_command(command: Commands) -> Result<()> {
    match command {
        Commands::Info { input } => cmd_info(input),
        Commands::Extract {
            input,
            output,
            start,
            end,
            frames,
            step,
        } => cmd_extract(input, output, start, end, frames, step),
        Commands::Convert {
            input,
            output,
            compression,
            full_precision,
            lossless,
            full_frame_interval,
        } => cmd_convert(
            input,
            output,
            compression,
            full_precision,
            lossless,
            full_frame_interval,
        ),
        Commands::Check { input, verbose } => cmd_check(input, verbose),
        Commands::Stats { input, detailed } => cmd_stats(input, detailed),
    }
}

/// Run CLI with provided arguments
pub fn run_cli_with_args(args: Vec<String>) -> Result<()> {
    match Cli::try_parse_from(args) {
        Ok(cli) => execute_command(cli.command),
        Err(err) => {
            // Check if this is a help or version request (which are not actual errors)
            match err.kind() {
                clap::error::ErrorKind::DisplayHelp | clap::error::ErrorKind::DisplayVersion => {
                    // Print the help/version output and return success
                    print!("{err}");
                    Ok(())
                }
                _ => {
                    // This is an actual error, propagate it
                    Err(anyhow::anyhow!("{}", err))
                }
            }
        }
    }
}

fn cmd_info(input: PathBuf) -> Result<()> {
    println!("Opening TACO file: {}", input.display());

    let mut reader = Reader::open(&input)?;
    let header = reader.header();

    println!("\n=== TACO File Information ===");
    println!("Format version: {}", header.version);
    println!("Number of atoms: {}", header.num_atoms);
    println!("Number of frames: {}", header.num_frames);
    println!("Time step: {} ps", header.time_step);

    println!("\n=== Simulation Metadata ===");
    let sim_meta = &header.simulation_metadata;
    if let Some(temp) = sim_meta.temperature {
        println!("Temperature: {temp} K");
    }
    if let Some(pressure) = sim_meta.pressure {
        println!("Pressure: {pressure} bar");
    }
    if let Some(ensemble) = &sim_meta.ensemble
        && !ensemble.is_empty()
    {
        println!("Ensemble: {ensemble}");
    }
    if let Some(software) = &sim_meta.software
        && !software.is_empty()
    {
        println!("Software: {software}");
    }

    println!("\n=== Compression Settings ===");
    let comp = &header.compression_settings;
    println!("Compression level: {}", comp.zstd_level);
    println!("Precision mode: {:?}", comp.precision);

    println!("\n=== Atom Metadata ===");
    let atom_meta = &header.atom_metadata;
    if let Some(elements) = &atom_meta.elements {
        println!("Number of elements: {}", elements.len());
        if !elements.is_empty() {
            let unique_elements: std::collections::HashSet<_> = elements.iter().collect();
            println!("Unique elements: {unique_elements:?}");
        }
    } else {
        println!("No element information available");
    }

    // File size
    if let Ok(metadata) = std::fs::metadata(&input) {
        let size_mb = metadata.len() as f64 / 1_048_576.0;
        println!("\n=== File Information ===");
        println!("File size: {size_mb:.2} MB");
        if header.num_frames > 0 {
            let mb_per_frame = size_mb / header.num_frames as f64;
            println!("Size per frame: {mb_per_frame:.3} MB");
        }
    }

    // Sample frame information (first few frames)
    println!("\n=== Sample Frame Information ===");
    let num_sample_frames = if header.num_frames <= 10 {
        header.num_frames // Show all frames if 10 or fewer
    } else {
        std::cmp::min(header.num_frames, 5) // Show first 5 if more than 10
    };

    if num_sample_frames > 0 {
        if header.num_frames <= 10 {
            println!("Showing all {num_sample_frames} frame(s):");
        } else {
            println!("Showing first {num_sample_frames} frame(s):");
        }
        for i in 0..num_sample_frames {
            match reader.read_frame(i) {
                Ok(frame) => {
                    print!("Frame {}: time={:.3}ps", i, frame.time);
                    if let Some(temp) = frame.temperature {
                        print!(", T={temp:.1}K");
                    } else {
                        print!(", T=None");
                    }
                    if let Some(pressure) = frame.pressure {
                        print!(", P={pressure:.1}bar");
                    } else {
                        print!(", P=None");
                    }
                    if let Some(pe) = frame.potential_energy {
                        print!(", PE={pe:.1}kJ/mol");
                    }
                    if let Some(ke) = frame.kinetic_energy {
                        print!(", KE={ke:.1}kJ/mol");
                    }
                    println!();
                }
                Err(e) => {
                    println!("Frame {i}: Error reading frame: {e}");
                    break;
                }
            }
        }
    } else {
        println!("No frames in file");
    }

    Ok(())
}

fn cmd_extract(
    input: PathBuf,
    output: PathBuf,
    start: Option<u64>,
    end: Option<u64>,
    frames: Option<Vec<u64>>,
    step: Option<u64>,
) -> Result<()> {
    println!("Extracting frames from: {}", input.display());
    println!("Output file: {}", output.display());

    let mut reader = Reader::open(&input)?;
    let header = reader.header().clone();
    let total_frames = header.num_frames;

    // Determine which frames to extract
    let frame_indices = if let Some(specific_frames) = frames {
        println!("Extracting specific frames: {specific_frames:?}");
        specific_frames
    } else {
        let start_frame = start.unwrap_or(0);
        let end_frame = end.unwrap_or(total_frames);
        let step_size = step.unwrap_or(1);
        if step_size == 0 {
            return Err(anyhow::anyhow!("Step size must be greater than 0"));
        }

        if start_frame >= total_frames {
            return Err(anyhow::anyhow!(
                "Start frame {} exceeds total frames {}",
                start_frame,
                total_frames
            ));
        }

        let actual_end = end_frame.min(total_frames);
        println!("Extracting frames {start_frame} to {actual_end} (step: {step_size})");

        (start_frame..actual_end)
            .step_by(step_size as usize)
            .collect()
    };

    if frame_indices.is_empty() {
        return Err(anyhow::anyhow!("No frames to extract"));
    }

    // Validate frame indices
    for &frame_idx in &frame_indices {
        if frame_idx >= total_frames {
            return Err(anyhow::anyhow!(
                "Frame index {} exceeds total frames {}",
                frame_idx,
                total_frames
            ));
        }
    }

    // Create a new writer with updated header
    let mut new_header = header.clone();
    new_header.num_frames = frame_indices.len() as u64;

    let mut writer = Writer::create_with_header(&output, new_header)?;

    // Read and write the selected frames
    print_info(&format!("Processing {} frames...", frame_indices.len()));
    for (i, &frame_idx) in frame_indices.iter().enumerate() {
        if i % 100 == 0 {
            print_info(&format!("Progress: {}/{}", i, frame_indices.len()));
        }

        let mut frame = reader.read_frame(frame_idx)?;
        // Update the frame number to be sequential
        frame.frame_number = i as u64;
        writer.write_frame(frame)?;
    }

    writer.finish()?;
    print_success(&format!(
        "Successfully extracted {} frames to {}",
        frame_indices.len(),
        output.display()
    ));

    Ok(())
}

fn cmd_convert(
    input: PathBuf,
    output: PathBuf,
    compression: i32,
    full_precision: bool,
    lossless: bool,
    full_frame_interval: u32,
) -> Result<()> {
    // For now, this assumes the input is already a TACO file
    // In the future, this could be extended to support other formats
    println!("Converting trajectory file: {}", input.display());
    println!("Output TACO file: {}", output.display());

    let mut reader = Reader::open(&input)?;
    let mut header = reader.header().clone();

    // Update compression settings
    let precision = if full_precision || lossless {
        PrecisionMode::Full
    } else {
        PrecisionMode::Half
    };

    header.compression_settings = CompressionSettings {
        precision,
        zstd_level: compression,
    };
    header.full_frame_interval = full_frame_interval;

    let mut writer = Writer::create_with_header(&output, header)?;

    let total_frames = reader.num_frames();
    println!("Converting {total_frames} frames with compression level {compression}...");

    // Copy all frames with new compression settings
    for i in 0..total_frames {
        if i % 100 == 0 {
            print_info(&format!("Progress: {i}/{total_frames}"));
        }

        let frame = reader.read_frame(i)?;
        writer.write_frame(frame)?;
    }

    writer.finish()?;
    print_success(&format!("Successfully converted to {}", output.display()));

    Ok(())
}

fn cmd_check(input: PathBuf, verbose: bool) -> Result<()> {
    println!("Checking TACO file: {}", input.display());

    let mut reader = Reader::open(&input)?;
    let total_frames = reader.num_frames();

    if verbose {
        println!("Validating {total_frames} frames...");
    }

    let mut errors = 0;

    for i in 0..total_frames {
        if verbose && i % 1000 == 0 {
            print_info(&format!("Checking frame {i}/{total_frames}"));
        }

        match reader.read_frame(i) {
            Ok(frame) => {
                // Basic validation
                if let Some(positions) = &frame.data.positions {
                    if positions.nrows() != reader.num_atoms() as usize {
                        print_error(&format!(
                            "Frame {i}: Position array has wrong number of atoms"
                        ));
                        errors += 1;
                    }

                    // Check for NaN or infinite values
                    for pos in positions.iter() {
                        if !pos.is_finite() {
                            print_error(&format!("Frame {i}: Invalid position value: {pos}"));
                            errors += 1;
                            break;
                        }
                    }
                }
            }
            Err(e) => {
                print_error(&format!("Frame {i}: Failed to read - {e}"));
                errors += 1;
            }
        }

        if errors > 10 {
            print_error("Too many errors, stopping validation");
            break;
        }
    }

    if errors == 0 {
        println!("✓ File validation successful - no errors found");
    } else {
        println!("✗ File validation failed - {errors} errors found");
    }

    Ok(())
}

fn cmd_stats(input: PathBuf, detailed: bool) -> Result<()> {
    println!("Analyzing TACO file: {}", input.display());

    let mut reader = Reader::open(&input)?;
    let total_frames = reader.num_frames();

    println!("\n=== Trajectory Statistics ===");

    if total_frames == 0 {
        println!("No frames to analyze");
        return Ok(());
    }

    // Sample frames for statistics (every 10th frame or max 1000 frames)
    let sample_step = (total_frames / 1000).clamp(1, 10);
    let sample_frames: Vec<u64> = (0..total_frames).step_by(sample_step as usize).collect();

    println!(
        "Analyzing {} sample frames (every {} frames)...",
        sample_frames.len(),
        sample_step
    );

    let mut min_positions = Vec::new();
    let mut max_positions = Vec::new();
    let mut total_frames_with_velocities = 0;
    let mut total_frames_with_forces = 0;

    for &frame_idx in &sample_frames {
        match reader.read_frame(frame_idx) {
            Ok(frame) => {
                if let Some(positions) = &frame.data.positions {
                    if min_positions.is_empty() {
                        min_positions = positions.iter().cloned().collect();
                        max_positions = positions.iter().cloned().collect();
                    } else {
                        for (i, &pos) in positions.iter().enumerate() {
                            if i < min_positions.len() {
                                min_positions[i] = min_positions[i].min(pos);
                                max_positions[i] = max_positions[i].max(pos);
                            }
                        }
                    }
                }

                if frame.data.velocities.is_some() {
                    total_frames_with_velocities += 1;
                }

                if frame.data.forces.is_some() {
                    total_frames_with_forces += 1;
                }
            }
            Err(e) => {
                eprintln!("Warning: Could not read frame {frame_idx}: {e}");
            }
        }
    }

    if !min_positions.is_empty() {
        let min_x = min_positions
            .iter()
            .step_by(3)
            .fold(f32::INFINITY, |a, &b| a.min(b));
        let max_x = max_positions
            .iter()
            .step_by(3)
            .fold(f32::NEG_INFINITY, |a, &b| a.max(b));
        let min_y = min_positions
            .iter()
            .skip(1)
            .step_by(3)
            .fold(f32::INFINITY, |a, &b| a.min(b));
        let max_y = max_positions
            .iter()
            .skip(1)
            .step_by(3)
            .fold(f32::NEG_INFINITY, |a, &b| a.max(b));
        let min_z = min_positions
            .iter()
            .skip(2)
            .step_by(3)
            .fold(f32::INFINITY, |a, &b| a.min(b));
        let max_z = max_positions
            .iter()
            .skip(2)
            .step_by(3)
            .fold(f32::NEG_INFINITY, |a, &b| a.max(b));

        println!("Position ranges:");
        println!(
            "  X: {:.3} to {:.3} Å (span: {:.3} Å)",
            min_x,
            max_x,
            max_x - min_x
        );
        println!(
            "  Y: {:.3} to {:.3} Å (span: {:.3} Å)",
            min_y,
            max_y,
            max_y - min_y
        );
        println!(
            "  Z: {:.3} to {:.3} Å (span: {:.3} Å)",
            min_z,
            max_z,
            max_z - min_z
        );
    }

    println!("Data availability:");
    println!(
        "  Frames with velocities: {}/{} ({:.1}%)",
        total_frames_with_velocities,
        sample_frames.len(),
        100.0 * total_frames_with_velocities as f64 / sample_frames.len() as f64
    );
    println!(
        "  Frames with forces: {}/{} ({:.1}%)",
        total_frames_with_forces,
        sample_frames.len(),
        100.0 * total_frames_with_forces as f64 / sample_frames.len() as f64
    );

    if detailed {
        println!("\n=== Detailed Frame Information ===");
        for i in 0..total_frames.min(20) {
            match reader.read_frame(i) {
                Ok(frame) => {
                    println!(
                        "Frame {}: time={:.3} ps, positions={}, velocities={}, forces={}",
                        i,
                        frame.time,
                        frame.data.positions.is_some(),
                        frame.data.velocities.is_some(),
                        frame.data.forces.is_some()
                    );
                }
                Err(e) => {
                    println!("Frame {i}: Error - {e}");
                }
            }
        }

        if total_frames > 20 {
            println!("... ({} more frames)", total_frames - 20);
        }
    }

    Ok(())
}
