use anyhow::{Context, Result, bail};
use chrono::{DateTime, Utc};
use dialoguer::Select;
use mithril_client::models::InstanceModel;
use mithril_client::models::bid_model::Status as BidStatus;
use mithril_client::models::instance_model::Status as InstanceStatus;
use std::collections::HashSet;
use std::os::unix::process::CommandExt;

use crate::api;
use crate::util::{WaitResult, calculate_age, wait_for_bid_instances, wait_for_ssh_ready};

pub async fn run(
    bid_name: Option<String>,
    node: Option<i32>,
    show: bool,
    no_wait: bool,
    command: Vec<String>,
) -> Result<()> {
    let client = api::Client::load()?;

    if let Some(ref bid_name) = bid_name {
        let (mut instances, bid, _project_fid) =
            client.find_instances_by_bid_name(bid_name).await?;

        // Check if we need to wait - instances must be running AND have SSH destinations
        let has_ready_instance = instances.iter().any(|i| {
            i.status == InstanceStatus::StatusRunning
                && i.ssh_destination.as_ref().is_some_and(|s| !s.is_empty())
        });
        let needs_wait = instances.is_empty() || !has_ready_instance;

        if needs_wait && !no_wait {
            // Check if bid is in a state where waiting makes sense
            match bid.status {
                BidStatus::Terminated => {
                    bail!("Bid '{}' is terminated - no instances available", bid.name);
                }
                BidStatus::Open | BidStatus::Allocated | BidStatus::Paused => {
                    match wait_for_bid_instances(&client, &bid.fid, &bid.name, None).await? {
                        WaitResult::Ready => {
                            // Re-fetch instances after waiting
                            let (new_instances, _, _) =
                                client.find_instances_by_bid_name(bid_name).await?;
                            instances = new_instances;
                        }
                        WaitResult::Terminated => {
                            bail!("Bid '{}' was terminated while waiting", bid.name);
                        }
                    }
                }
                BidStatus::Preempting => {
                    bail!(
                        "Bid '{}' is being preempted - instances are shutting down",
                        bid.name
                    );
                }
            }
        } else if needs_wait && no_wait {
            bail!(
                "No running instances for bid '{}'. Use without --no-wait to wait for instances.",
                bid.name
            );
        }

        if let Some(n) = node
            && (n < 0 || (n as usize) >= instances.len())
        {
            bail!(
                "Node index {} out of range. Bid '{}' has {} instance(s) (valid: 0-{}).",
                n,
                bid.name,
                instances.len(),
                instances.len().saturating_sub(1)
            );
        }

        if command.is_empty() {
            let idx = node.unwrap_or(0) as usize;
            let instance = &instances[idx];

            // Wait for SSH to be ready if not in no-wait mode
            if !no_wait
                && !show
                && let Some(ref dest) = instance.ssh_destination
            {
                wait_for_ssh_ready(dest).await?;
            }

            exec_ssh_interactive(instance, idx, instances.len(), show)?;
        } else if let Some(n) = node {
            let instance = &instances[n as usize];

            // Wait for SSH to be ready if not in no-wait mode
            if !no_wait
                && !show
                && let Some(ref dest) = instance.ssh_destination
            {
                wait_for_ssh_ready(dest).await?;
            }

            exec_ssh_single_command(instance, n as usize, &command, show)?;
        } else {
            // For multi-node commands, wait for all SSH connections if not in no-wait mode
            if !no_wait && !show {
                for inst in &instances {
                    if let Some(ref dest) = inst.ssh_destination {
                        wait_for_ssh_ready(dest).await?;
                    }
                }
            }

            exec_ssh_all_nodes(&instances, &command, show).await?;
        }
    } else {
        let (instance, _project_fid) = select_instance_interactive(&client).await?;
        exec_ssh_interactive(&instance, 0, 1, show)?;
    }

    Ok(())
}

fn exec_ssh_interactive(
    instance: &InstanceModel,
    node_idx: usize,
    total_nodes: usize,
    show: bool,
) -> Result<()> {
    let ssh_destination = instance
        .ssh_destination
        .as_ref()
        .context("Instance does not have an SSH destination yet. It may still be provisioning.")?;

    let ssh_target = format!("ubuntu@{ssh_destination}");

    if show {
        if total_nodes > 1 {
            println!("[node-{node_idx}] ssh {ssh_target}");
        } else {
            println!("ssh {ssh_target}");
        }
        return Ok(());
    }

    if total_nodes > 1 {
        println!("Connecting to {ssh_target} (node {node_idx}/{total_nodes})...");
    } else {
        println!("Connecting to {ssh_target}...");
    }

    let mut cmd = std::process::Command::new("ssh");
    cmd.arg(&ssh_target);
    let err = cmd.exec();
    bail!("Failed to execute ssh: {err}");
}

fn exec_ssh_single_command(
    instance: &InstanceModel,
    node_idx: usize,
    command: &[String],
    show: bool,
) -> Result<()> {
    let ssh_destination = instance
        .ssh_destination
        .as_ref()
        .context("Instance does not have an SSH destination yet. It may still be provisioning.")?;

    let ssh_target = format!("ubuntu@{ssh_destination}");

    if show {
        println!(
            "[node-{}] ssh {} {}",
            node_idx,
            ssh_target,
            command.join(" ")
        );
        return Ok(());
    }

    println!("[node-{}] Running: {}", node_idx, command.join(" "));

    let mut cmd = std::process::Command::new("ssh");
    cmd.arg(&ssh_target).args(command);
    let err = cmd.exec();
    bail!("Failed to execute ssh: {err}");
}

async fn exec_ssh_all_nodes(
    instances: &[InstanceModel],
    command: &[String],
    show: bool,
) -> Result<()> {
    if show {
        for (i, inst) in instances.iter().enumerate() {
            if let Some(ref dest) = inst.ssh_destination {
                println!("[node-{}] ssh ubuntu@{} {}", i, dest, command.join(" "));
            } else {
                println!("[node-{i}] (no SSH destination available)");
            }
        }
        return Ok(());
    }

    println!(
        "Running command on {} node(s) in parallel: {}",
        instances.len(),
        command.join(" ")
    );

    let handles: Vec<_> = instances
        .iter()
        .enumerate()
        .filter_map(|(i, inst)| {
            let ssh_dest = inst.ssh_destination.clone()?;
            let cmd = command.to_vec();
            Some(tokio::spawn(async move {
                let output = tokio::process::Command::new("ssh")
                    .arg(format!("ubuntu@{ssh_dest}"))
                    .args(&cmd)
                    .output()
                    .await?;
                Ok::<_, anyhow::Error>((i, output))
            }))
        })
        .collect();

    let results = futures::future::join_all(handles).await;

    for result in results {
        let (node_idx, output) = result.context("Task join error")??;
        for line in String::from_utf8_lossy(&output.stdout).lines() {
            println!("[node-{node_idx}] {line}");
        }
        for line in String::from_utf8_lossy(&output.stderr).lines() {
            eprintln!("[node-{node_idx}] {line}");
        }
    }

    Ok(())
}

async fn select_instance_interactive(client: &api::Client) -> Result<(InstanceModel, String)> {
    println!("Fetching all running instances...");

    let projects = client.fetch_projects().await?;
    println!("Found projects: {projects:?}");

    let instance_types = client.fetch_instance_types().await?;

    let mut all_fetched: Vec<(InstanceModel, String)> = Vec::new();
    for project in &projects {
        println!("Fetching instances for project: {}", project.fid);
        let instances = client.fetch_instances_for_project(&project.fid).await?;
        all_fetched.extend(instances);
    }

    println!(
        "Successfully fetched instances from {} projects",
        projects.len()
    );

    let mut filtered_statuses: HashSet<InstanceStatus> = HashSet::new();

    let mut all_instances: Vec<(InstanceModel, String, bool)> = all_fetched
        .into_iter()
        .filter_map(|(inst, proj)| {
            let is_selectable = inst.status == InstanceStatus::StatusRunning;
            let is_visible = is_selectable
                || (inst.status != InstanceStatus::StatusTerminated
                    && inst.status != InstanceStatus::StatusError);

            if is_visible {
                Some((inst, proj, is_selectable))
            } else {
                filtered_statuses.insert(inst.status);
                None
            }
        })
        .collect();

    if !filtered_statuses.is_empty() {
        let mut statuses_vec: Vec<_> = filtered_statuses.into_iter().collect();
        statuses_vec.sort();
        let status_names: Vec<String> = statuses_vec.iter().map(|s| format!("{s:?}")).collect();
        log::debug!(
            "Filtered out instances with statuses: {}",
            status_names.join(", ")
        );
    }

    if all_instances.is_empty() {
        bail!("No instances found");
    }

    all_instances.sort_by(|a, b| b.0.created_at.cmp(&a.0.created_at));

    let selectable_count = all_instances.iter().filter(|(_, _, sel)| *sel).count();

    if selectable_count == 1 && all_instances.len() == 1 {
        let (instance, project_fid, _) = all_instances.into_iter().next().unwrap();
        println!("Found 1 running instance: {}", instance.name);
        return Ok((instance, project_fid));
    }

    println!(
        "Found {} instances ({} running)",
        all_instances.len(),
        selectable_count
    );

    let now = Utc::now();
    let one_week_ago = now - chrono::Duration::days(7);

    let old_unselectable: Vec<usize> = all_instances
        .iter()
        .enumerate()
        .filter_map(|(idx, (inst, _, selectable))| {
            if !selectable
                && let Ok(created) = DateTime::parse_from_rfc3339(&inst.created_at)
                && created < one_week_ago
            {
                return Some(idx);
            }
            None
        })
        .collect();

    let mut show_old = false;

    loop {
        let visible_instances: Vec<usize> = all_instances
            .iter()
            .enumerate()
            .filter_map(|(idx, _)| {
                if show_old || !old_unselectable.contains(&idx) {
                    Some(idx)
                } else {
                    None
                }
            })
            .collect();

        let mut display_items: Vec<String> = visible_instances
            .iter()
            .map(|&idx| {
                let (inst, _proj, selectable) = &all_instances[idx];
                let age = calculate_age(&inst.created_at);
                let instance_type_info = instance_types.get(&inst.instance_type);

                let gpu_info = if let Some(it) = instance_type_info {
                    if it.num_gpus > 0 {
                        format!(" | {}x {}", it.num_gpus, it.gpu_type)
                    } else {
                        String::new()
                    }
                } else {
                    String::new()
                };

                let base_text = format!(
                    "{} | {}{} | {} | {:?} | {} | {}",
                    inst.name,
                    inst.instance_type,
                    gpu_info,
                    inst.region.as_deref().unwrap_or("unknown"),
                    inst.status,
                    inst.created_by,
                    age
                );

                if !selectable {
                    format!("\x1b[31m{base_text} [unavailable]\x1b[0m")
                } else {
                    base_text
                }
            })
            .collect();

        if !old_unselectable.is_empty() {
            if show_old {
                display_items.push(format!(
                    "\x1b[90m--- Hide {} old unavailable instances ---\x1b[0m",
                    old_unselectable.len()
                ));
            } else {
                display_items.push(format!(
                    "\x1b[90m--- Show {} more old unavailable instances ---\x1b[0m",
                    old_unselectable.len()
                ));
            }
        }

        let selection = Select::new()
            .with_prompt("Select instance to SSH into (unavailable instances cannot be selected)")
            .items(&display_items)
            .default(0)
            .interact()
            .context("Failed to get user selection")?;

        if selection == display_items.len() - 1 && !old_unselectable.is_empty() {
            show_old = !show_old;
            continue;
        }

        let actual_idx = visible_instances[selection];
        let (instance, project_fid, selectable) = &all_instances[actual_idx];

        if *selectable {
            return Ok((instance.clone(), project_fid.clone()));
        } else {
            println!(
                "⚠ Cannot SSH into instance '{}' with status '{:?}'. Please select a running instance.",
                instance.name, instance.status
            );
        }
    }
}
