molx 0.1.0

use std::fs::File;
use std::io::{self, BufWriter, IsTerminal};
use std::path::PathBuf;

use clap::{Args, Parser, Subcommand, ValueEnum};
use molx::atlas::{AtlasClient, read_sequence_file, validate_atlas_id, write_pdb};
use molx::{ColorScheme, Protein, Representation, Viewer};
use rstui::{Runner, RunnerConfig};

#[derive(Parser)]
#[command(
    name = "molx",
    version,
    about = "Explore protein structures in a graphical terminal",
    long_about = None
)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Open a local PDB structure in the interactive 3D viewer.
    View(ViewArgs),
    /// Print a compact summary of a PDB structure.
    Info {
        /// Local PDB file.
        path: PathBuf,
    },
    /// Render a PDB structure to a PNG without entering the interactive viewer.
    Snapshot(SnapshotArgs),
    /// Fetch predicted structures or run a small ESMFold job.
    Atlas {
        #[command(subcommand)]
        command: AtlasCommand,
    },
}

#[derive(Args)]
struct ViewArgs {
    /// Local PDB file.
    path: PathBuf,
    /// Initial molecular representation.
    #[arg(long, value_enum, default_value_t = RepresentationArg::Backbone)]
    representation: RepresentationArg,
    /// Initial color scheme; auto selects pLDDT for ESMFold structures.
    #[arg(long, value_enum, default_value_t = ColorArg::Auto)]
    color: ColorArg,
}

#[derive(Args)]
struct SnapshotArgs {
    /// Local PDB file.
    path: PathBuf,
    /// Destination PNG file.
    #[arg(short, long, default_value = "molx.png")]
    output: PathBuf,
    /// Image width in pixels.
    #[arg(long, default_value_t = 1200)]
    width: u32,
    /// Image height in pixels.
    #[arg(long, default_value_t = 800)]
    height: u32,
    /// Molecular representation.
    #[arg(long, value_enum, default_value_t = RepresentationArg::Backbone)]
    representation: RepresentationArg,
    /// Color scheme; auto selects pLDDT for ESMFold structures.
    #[arg(long, value_enum, default_value_t = ColorArg::Auto)]
    color: ColorArg,
    /// Replace an existing PNG file.
    #[arg(long)]
    force: bool,
}

#[derive(Subcommand)]
enum AtlasCommand {
    /// Download a predicted PDB structure by MGnify/ESM Atlas ID.
    Fetch {
        /// Identifier beginning with MGYP, for example MGYP002537940442.
        id: String,
        /// Destination PDB path; defaults to <ID>.pdb.
        #[arg(short, long)]
        output: Option<PathBuf>,
        /// Open the downloaded structure after saving it.
        #[arg(long)]
        view: bool,
        /// Replace an existing output file.
        #[arg(long)]
        force: bool,
    },
    /// Predict a structure from one protein sequence with the public ESMFold API.
    Fold(FoldArgs),
}

#[derive(Args)]
struct FoldArgs {
    /// Amino-acid sequence as a command-line string.
    #[arg(long, required_unless_present = "fasta", conflicts_with = "fasta")]
    sequence: Option<String>,
    /// FASTA or plain sequence file.
    #[arg(
        long,
        required_unless_present = "sequence",
        conflicts_with = "sequence"
    )]
    fasta: Option<PathBuf>,
    /// Destination PDB path.
    #[arg(short, long, default_value = "esmfold.pdb")]
    output: PathBuf,
    /// Open the prediction after saving it.
    #[arg(long)]
    view: bool,
    /// Replace an existing output file.
    #[arg(long)]
    force: bool,
}

#[derive(Clone, Copy, Debug, ValueEnum)]
enum RepresentationArg {
    Backbone,
    Atoms,
    Combined,
}

impl From<RepresentationArg> for Representation {
    fn from(value: RepresentationArg) -> Self {
        match value {
            RepresentationArg::Backbone => Self::Backbone,
            RepresentationArg::Atoms => Self::Atoms,
            RepresentationArg::Combined => Self::Combined,
        }
    }
}

#[derive(Clone, Copy, Debug, ValueEnum)]
enum ColorArg {
    Auto,
    Chain,
    Element,
    Confidence,
}

fn main() {
    if let Err(error) = run(Cli::parse()) {
        eprintln!("molx: {error}");
        std::process::exit(1);
    }
}

fn run(cli: Cli) -> Result<(), String> {
    match cli.command {
        Command::View(arguments) => view_path(
            arguments.path,
            arguments.representation.into(),
            arguments.color,
        ),
        Command::Info { path } => {
            println!("{}", Protein::from_path(path)?.summary());
            Ok(())
        }
        Command::Snapshot(arguments) => snapshot(arguments),
        Command::Atlas { command } => run_atlas(command),
    }
}

fn snapshot(arguments: SnapshotArgs) -> Result<(), String> {
    if arguments.output.exists() && !arguments.force {
        return Err(format!(
            "{} already exists; pass --force to replace it",
            arguments.output.display()
        ));
    }
    let protein = Protein::from_path(arguments.path)?;
    let color = resolve_color(arguments.color, &protein);
    let viewer = Viewer::new(protein, arguments.representation.into(), color);
    let frame = viewer.snapshot(arguments.width, arguments.height);
    let file = File::create(&arguments.output)
        .map_err(|error| format!("cannot create {}: {error}", arguments.output.display()))?;
    let mut encoder = png::Encoder::new(BufWriter::new(file), frame.width, frame.height);
    encoder.set_color(png::ColorType::Rgba);
    encoder.set_depth(png::BitDepth::Eight);
    let mut writer = encoder
        .write_header()
        .map_err(|error| format!("cannot initialize PNG output: {error}"))?;
    writer
        .write_image_data(&frame.pixels)
        .map_err(|error| format!("cannot write {}: {error}", arguments.output.display()))?;
    println!("Saved {}", arguments.output.display());
    Ok(())
}

fn run_atlas(command: AtlasCommand) -> Result<(), String> {
    match command {
        AtlasCommand::Fetch {
            id,
            output,
            view,
            force,
        } => {
            let id = validate_atlas_id(&id)?;
            let output = output.unwrap_or_else(|| PathBuf::from(format!("{id}.pdb")));
            eprintln!("Fetching {id} from ESM Atlas...");
            let pdb = AtlasClient::new()?.fetch_structure(&id)?;
            write_pdb(&output, &pdb, force)?;
            eprintln!("Saved {}", output.display());
            if view {
                view_path(output, Representation::Backbone, ColorArg::Confidence)
            } else {
                Ok(())
            }
        }
        AtlasCommand::Fold(arguments) => {
            let sequence = if let Some(sequence) = arguments.sequence {
                molx::atlas::normalize_sequence(&sequence)?
            } else if let Some(path) = arguments.fasta {
                read_sequence_file(path)?
            } else {
                unreachable!("clap requires one sequence input")
            };
            eprintln!(
                "Folding {} residues with the public ESMFold service...",
                sequence.len()
            );
            let pdb = AtlasClient::new()?.fold_sequence(&sequence)?;
            write_pdb(&arguments.output, &pdb, arguments.force)?;
            eprintln!("Saved {}", arguments.output.display());
            if arguments.view {
                view_path(
                    arguments.output,
                    Representation::Backbone,
                    ColorArg::Confidence,
                )
            } else {
                Ok(())
            }
        }
    }
}

fn view_path(
    path: PathBuf,
    representation: Representation,
    color_argument: ColorArg,
) -> Result<(), String> {
    if !io::stdout().is_terminal() {
        return Err("the interactive viewer needs a terminal on stdout; use `molx info` for non-interactive output".to_owned());
    }
    let protein = Protein::from_path(path)?;
    let color = resolve_color(color_argument, &protein);
    let viewer = Viewer::new(protein, representation, color);
    Runner::new(RunnerConfig {
        frame_rate: 30,
        idle_refine_ms: 110,
        ..RunnerConfig::default()
    })
    .run(viewer)
}

fn resolve_color(color_argument: ColorArg, protein: &Protein) -> ColorScheme {
    match color_argument {
        ColorArg::Auto if protein.is_prediction => ColorScheme::Confidence,
        ColorArg::Auto | ColorArg::Chain => ColorScheme::Chain,
        ColorArg::Element => ColorScheme::Element,
        ColorArg::Confidence => ColorScheme::Confidence,
    }
}