prox 0.1.1

use crate::{
    Config, Proc, ProcStatus, Prox, ProxEvent, ProxSignal, debug, error, info,
    logging::{self, DEFAULT_PROX_PREFIX, Level},
    warn,
};
use anyhow::{Context, Result};
use dashmap::DashMap;
use notify::{RecursiveMode, Watcher, recommended_watcher};
use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    process::{Child, Command},
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
        mpsc,
    },
    thread::{sleep, spawn},
    time::{Duration, Instant},
};

impl Prox {
    /// Get a reference to the configuration
    pub fn config(&self) -> &Config {
        &self.config
    }

    /// Get a reference to the processes
    pub fn procs(&self) -> &[Proc] {
        &self.procs
    }

    /// Set up a [`ProxEvent`] receiver channel and return the receiver.
    /// Can be run only once. Returns error if run again.
    pub fn setup_event_rx(&mut self) -> Result<mpsc::Receiver<ProxEvent>> {
        if self.event_tx.is_some() {
            return Err(anyhow::anyhow!("event channel is already set up"));
        }
        let (tx, rx) = mpsc::channel();
        self.event_tx = Some(tx);
        Ok(rx)
    }

    /// Start all processes
    pub fn start(&mut self) -> Result<()> {
        // Initialize logging
        logging::init_logging(logging::LogConfig {
            prox_prefix: DEFAULT_PROX_PREFIX.to_string(),
            prefix_width: self.config.prefix_width,
            timestamp: self.config.show_timestamps,
            level: Level::from_env(),
        });

        // TODO: debounce these to public ProxEvent::Idle
        // use self.config.output_idle_debounce_ms
        let (proc_output_tx, proc_output_rx) = mpsc::channel::<()>();
        let (restart_tx, restart_rx) = mpsc::channel::<String>();

        // Set up idle detection if enabled
        if self.config.output_idle_debounce_ms > 0 {
            let idle_debounce_duration = Duration::from_millis(self.config.output_idle_debounce_ms);
            let event_tx = self.event_tx.clone();
            let running = self.running.clone();

            spawn(move || {
                let mut idle_sent = false;

                loop {
                    if !running.load(Ordering::SeqCst) {
                        break;
                    }

                    match proc_output_rx.recv_timeout(idle_debounce_duration) {
                        Ok(()) => {
                            // Received output, reset idle state
                            idle_sent = false;
                        }
                        Err(mpsc::RecvTimeoutError::Timeout) => {
                            // Timeout reached, send idle event if not already sent
                            if !idle_sent {
                                if let Some(ref event_tx) = event_tx {
                                    event_tx.send(ProxEvent::Idle).ok();
                                }
                                idle_sent = true;
                            }
                        }
                        Err(mpsc::RecvTimeoutError::Disconnected) => {
                            // Channel disconnected, exit
                            break;
                        }
                    }
                }
            });
        }

        self.running.store(true, Ordering::SeqCst);
        let keep_going = self.config.keep_going;

        self.assign_missing_colors();

        if let Some(rx) = self.signal_rx.take() {
            let running = self.running.clone();
            let child_refs = self.child_refs.clone();
            let restart_tx = restart_tx.clone();
            let procs = self.procs.clone();
            spawn(move || {
                for signal in rx {
                    match signal {
                        ProxSignal::Start => {
                            // restart any processes that are not currently running
                            let mut start_count = 0;
                            for proc in &procs {
                                if let Some(mut child) = child_refs.get_mut(&proc.name) {
                                    if child.try_wait().is_ok_and(|s| s.is_none()) {
                                        continue;
                                    }
                                }
                                start_count += 1;
                                restart_tx.send(proc.name.clone()).ok();
                            }
                            if start_count == 0 {
                                info!("All processes are already running");
                            }
                        }
                        ProxSignal::Restart => {
                            error!("RECEIVED RESTART SIGNAL");
                            // restart all processes
                            for proc in &procs {
                                restart_tx.send(proc.name.clone()).ok();
                            }
                        }
                        ProxSignal::Shutdown => {
                            running.store(false, Ordering::SeqCst);
                            return;
                        }
                    }
                }
            });
        }

        // Set up Control-C handler
        if self.config.handle_control_c {
            let running = self.running.clone();
            let event_tx = self.event_tx.clone();
            if let Err(err) = ctrlc::try_set_handler(move || {
                running.store(false, Ordering::SeqCst);
                if let Some(event_tx) = &event_tx {
                    event_tx.send(ProxEvent::SigIntReceived).ok();
                }
            }) {
                warn!("Failed to set Control-C handler (is one already set?): {err}");
            }
        }

        // Start all processes
        let (failure_tx, failure_rx) = mpsc::channel();
        for proc in &self.procs {
            if let Err(err) =
                self.start_process(proc.clone(), failure_tx.clone(), proc_output_tx.clone())
            {
                if !keep_going {
                    eprintln!("Failed to start process {}: {err}", proc.name);
                    self.shutdown();
                    return Err(err);
                }
            }
        }

        // Set up file watchers
        if let Err(err) = self.setup_file_watchers(failure_tx.clone(), proc_output_tx.clone()) {
            eprintln!("Failed to set up file watchers: {err}");
            self.shutdown();
            return Err(err);
        }

        // Main loop
        while self.running.load(Ordering::SeqCst) {
            // Check for restart requests
            if let Ok(proc_name) = restart_rx.try_recv()
                && let Some(proc) = self.procs.iter().find(|p| p.name == proc_name)
                && let Err(err) =
                    self.start_process(proc.clone(), failure_tx.clone(), proc_output_tx.clone())
            {
                eprintln!("Failed to restart process {}: {err}", proc.name);
                if !keep_going {
                    self.running.store(false, Ordering::SeqCst);
                }
                break;
            }

            // Check for failures
            if let Ok(failure_msg) = failure_rx.try_recv() {
                if !keep_going {
                    eprintln!("Process failure: {failure_msg}");
                    self.running.store(false, Ordering::SeqCst);
                    break;
                }
            }

            // Check if any processes have exited
            let mut any_exited = false;
            let mut dead_processes = Vec::new();

            for mut entry in self.child_refs.iter_mut() {
                let name = entry.key().clone();
                let child = entry.value_mut();

                match child.try_wait() {
                    Ok(Some(status)) => {
                        println!("{name} exited with status: {status:?}");
                        if let Some(event_tx) = &self.event_tx {
                            event_tx
                                .send(ProxEvent::Exited {
                                    proc_name: name.clone(),
                                    status,
                                })
                                .ok();
                        }
                        if !keep_going {
                            self.running.store(false, Ordering::SeqCst);
                            any_exited = true;
                            break;
                        } else {
                            println!("{name} will restart when watched files change");
                            dead_processes.push(name);
                        }
                    }
                    Ok(None) => {}
                    Err(e) => {
                        println!("Failed to wait on process {name}: {e}");
                        if !keep_going {
                            self.running.store(false, Ordering::SeqCst);
                            any_exited = true;
                            break;
                        } else {
                            println!("{name} will restart when watched files change");
                            dead_processes.push(name);
                        }
                    }
                }
            }

            // Remove dead processes from child_refs when keep_going is true
            // Do this carefully to avoid interfering with restart logic
            for dead_process in dead_processes {
                // Only remove if not currently starting (to avoid race with restart)
                if !self.starting.contains_key(&dead_process) {
                    self.child_refs.remove(&dead_process);
                }
            }

            if any_exited {
                break;
            }

            // Sleep briefly
            sleep(Duration::from_millis(100));
        }

        error!("SHUTDOWN NOW");

        // Shutdown
        self.shutdown();
        Ok(())
    }

    fn assign_missing_colors(&mut self) {
        // Assign colors to processes that don't have one
        let colors = &self.config.colors;
        for (i, proc) in self.procs.iter_mut().enumerate() {
            proc.color.get_or_insert_with(|| colors[i % colors.len()]);
        }
    }

    fn start_process(
        &self,
        proc: Proc,
        failure_tx: mpsc::Sender<String>,
        proc_output_tx: mpsc::Sender<()>,
    ) -> Result<()> {
        let child_refs = self.child_refs.clone();
        let starting = self.starting.clone();
        let running = self.running.clone();
        let fallback_timeout = self.config.readiness_fallback_timeout;
        let process_group_id = self.process_group_id.clone();
        let config = self.config.clone();
        let procs = self.procs.clone();
        let event_tx = self.event_tx.clone();
        let status_refs = self.status_refs.clone();

        spawn(move || {
            starting.insert(proc.name.clone(), Instant::now());

            // Remove any existing process
            if let Some((_, mut old_child)) = child_refs.remove(&proc.name) {
                old_child.kill().ok();
                old_child.wait().ok();
            }

            status_refs.insert(proc.name.clone(), ProcStatus::Starting);

            match proc.start_service(&process_group_id, &config, Some(proc_output_tx.clone())) {
                Ok((new_child, readiness_rx)) => {
                    child_refs.insert(proc.name.clone(), new_child);
                    info!("{} Service started, waiting for readiness...", proc.name);

                    if proc.readiness_pattern.is_some() {
                        // Wait for readiness pattern (forever, until process crashes)
                        loop {
                            match readiness_rx.try_recv() {
                                Ok(()) => {
                                    status_refs.insert(proc.name.clone(), ProcStatus::Running);

                                    info!(
                                        "{} detected readiness pattern, start complete",
                                        proc.name
                                    );
                                    if let Some(event_tx) = &event_tx {
                                        event_tx
                                            .send(ProxEvent::Started {
                                                proc_name: proc.name.clone(),
                                            })
                                            .ok();
                                    }
                                    break;
                                }
                                Err(mpsc::TryRecvError::Empty) => {
                                    if !running.load(Ordering::SeqCst) {
                                        break;
                                    }
                                    sleep(Duration::from_millis(100));
                                    continue;
                                }
                                Err(mpsc::TryRecvError::Disconnected) => {
                                    // Process died, will be detected by crash detection
                                    break;
                                }
                            }
                        }
                    } else {
                        // No readiness pattern - wait fallback timeout then assume ready
                        let start_time = Instant::now();
                        while start_time.elapsed() < fallback_timeout
                            && running.load(Ordering::SeqCst)
                        {
                            // Check if process is still alive
                            if let Some(mut child_entry) = child_refs.get_mut(&proc.name) {
                                match child_entry.try_wait() {
                                    Ok(Some(status)) => {
                                        // Process exited, will be detected by crash detection
                                        status_refs
                                            .insert(proc.name.clone(), ProcStatus::Exited(status));

                                        break;
                                    }
                                    Ok(None) => {
                                        // Still running, continue waiting
                                    }
                                    Err(err) => {
                                        // Error checking status, assume dead
                                        status_refs.insert(
                                            proc.name.clone(),
                                            ProcStatus::Error(err.to_string()),
                                        );

                                        break;
                                    }
                                }
                            }
                            sleep(Duration::from_millis(100));
                        }

                        info!("{} assumed running after fallback timeout", proc.name);
                        status_refs.insert(proc.name.clone(), ProcStatus::Running);

                        if let Some(event_tx) = &event_tx {
                            event_tx
                                .send(ProxEvent::Started {
                                    proc_name: proc.name.clone(),
                                })
                                .ok();
                        }
                    }

                    starting.remove(&proc.name);

                    if starting.is_empty() {
                        let exited_procs = exited_procs(&procs, &child_refs);
                        if exited_procs.is_empty() {
                            info!("All processes started successfully!");
                            if let Some(event_tx) = &event_tx {
                                event_tx.send(ProxEvent::AllStarted).ok();
                            }
                        } else {
                            info!("Procs exited! {exited_procs:?}");
                            if let Some(event_tx) = &event_tx {
                                event_tx.send(ProxEvent::SomeFailed { exited_procs }).ok();
                            }
                        }
                    } else {
                        debug!(
                            "starting.is_empty(): {}, child_refs.len(): {}, total_procs: {}",
                            starting.is_empty(),
                            child_refs.len(),
                            procs.len()
                        );
                    }
                }
                Err(e) => {
                    let error_msg = format!("Failed to start {}: {e}", proc.name);

                    status_refs.insert(proc.name.clone(), ProcStatus::Error(error_msg.clone()));

                    if let Some(event_tx) = &event_tx {
                        event_tx
                            .send(ProxEvent::StartFailed {
                                proc_name: proc.name.clone(),
                                message: error_msg.clone(),
                            })
                            .ok();
                    }
                    error!("{error_msg}");
                    failure_tx.send(error_msg).ok();
                    starting.remove(&proc.name);
                    running.store(false, Ordering::SeqCst);
                }
            }
        });

        Ok(())
    }

    fn setup_file_watchers(
        &self,
        failure_tx: mpsc::Sender<String>,
        proc_output_tx: mpsc::Sender<()>,
    ) -> Result<()> {
        let mut watch_path_procs: HashMap<PathBuf, Vec<Proc>> = HashMap::new();

        // Build path -> process mapping
        for proc in &self.procs {
            for path in proc.watch_abs(&self.config)? {
                watch_path_procs.entry(path).or_default().push(proc.clone());
            }
        }

        if watch_path_procs.is_empty() {
            return Ok(());
        }

        // Start file watcher
        let running = self.running.clone();
        let child_refs = self.child_refs.clone();
        let starting = self.starting.clone();
        let fallback_timeout = self.config.readiness_fallback_timeout;
        let process_group_id = self.process_group_id.clone();
        let config = self.config.clone();
        let procs = self.procs.clone();
        let event_tx = self.event_tx.clone();
        let proc_output_tx_clone = proc_output_tx.clone();

        spawn(move || {
            let (tx, rx) = mpsc::channel();
            let mut watcher = match recommended_watcher(tx) {
                Ok(w) => w,
                Err(e) => {
                    eprintln!("Failed to create file watcher: {e}");
                    return;
                }
            };

            for path in watch_path_procs.keys() {
                println!("[GLOBAL] Watching: {}", path.display());
                if let Err(e) = watcher.watch(path, RecursiveMode::Recursive) {
                    eprintln!("Failed to watch path {}: {e}", path.display());
                }
            }

            while running.load(Ordering::SeqCst) {
                match rx.recv() {
                    Ok(Ok(event)) => {
                        let mut procs_to_restart = HashMap::new();

                        for event_path in &event.paths {
                            for (watched_path, procs) in &watch_path_procs {
                                if !event_path.starts_with(watched_path) {
                                    // println!("SKIP WATCH PATH: {}", event_path.display());
                                    continue;
                                }

                                println!("RESTART FOR WATCH PATH: {}", event_path.display());

                                for proc in procs {
                                    println!(
                                        "restart? {}, starting? {}",
                                        proc.name,
                                        starting.contains_key(&proc.name)
                                    );
                                    if !starting.contains_key(&proc.name) {
                                        println!("RESTART PROC: {}", proc.name);
                                        procs_to_restart.insert(proc.name.clone(), proc.clone());
                                    }
                                }
                            }
                        }

                        for (_, proc) in procs_to_restart {
                            let child_refs = child_refs.clone();
                            let starting = starting.clone();
                            starting.insert(proc.name.clone(), Instant::now());
                            let running = running.clone();
                            let failure_tx = failure_tx.clone();
                            let process_group_id = process_group_id.clone();
                            let config = config.clone();
                            let procs = procs.clone();
                            let event_tx = event_tx.clone();
                            let proc_output_tx = proc_output_tx_clone.clone();

                            spawn(move || {
                                // Small delay to avoid race conditions with process cleanup
                                sleep(Duration::from_millis(100));

                                // Remove existing process
                                if let Some((_, mut old_child)) = child_refs.remove(&proc.name) {
                                    old_child.kill().ok();
                                    old_child.wait().ok();
                                }

                                match proc.start_service(
                                    &process_group_id,
                                    &config,
                                    Some(proc_output_tx.clone()),
                                ) {
                                    Ok((new_child, readiness_rx)) => {
                                        child_refs.insert(proc.name.clone(), new_child);

                                        if proc.readiness_pattern.is_some() {
                                            // Wait for readiness pattern (forever, until process crashes)
                                            loop {
                                                match readiness_rx.try_recv() {
                                                    Ok(()) => {
                                                        println!(
                                                            "[GLOBAL] {} restart complete",
                                                            proc.name
                                                        );
                                                        if let Some(event_tx) = &event_tx {
                                                            event_tx
                                                                .send(ProxEvent::Restarted {
                                                                    proc_name: proc.name.clone(),
                                                                })
                                                                .ok();
                                                        }
                                                        break;
                                                    }
                                                    Err(mpsc::TryRecvError::Empty) => {
                                                        if !running.load(Ordering::SeqCst) {
                                                            break;
                                                        }
                                                        sleep(Duration::from_millis(10));
                                                        continue;
                                                    }
                                                    Err(mpsc::TryRecvError::Disconnected) => {
                                                        // Process died, will be detected by crash detection
                                                        break;
                                                    }
                                                }
                                            }
                                        } else {
                                            // No readiness pattern - wait fallback timeout then assume ready
                                            let start_time = Instant::now();
                                            while start_time.elapsed() < fallback_timeout
                                                && running.load(Ordering::SeqCst)
                                            {
                                                // Check if process is still alive
                                                if let Some(mut child_entry) =
                                                    child_refs.get_mut(&proc.name)
                                                {
                                                    match child_entry.try_wait() {
                                                        Ok(Some(_)) => {
                                                            // Process exited, will be detected by crash detection
                                                            break;
                                                        }
                                                        Ok(None) => {
                                                            // Still running, continue waiting
                                                        }
                                                        Err(_) => {
                                                            // Error checking status, assume dead
                                                            break;
                                                        }
                                                    }
                                                }
                                                sleep(Duration::from_millis(100));
                                            }

                                            // If we get here and process is still alive, assume ready
                                            println!(
                                                "[GLOBAL] {} restart assumed ready after fallback timeout",
                                                proc.name
                                            );
                                        }

                                        starting.remove(&proc.name);
                                        if starting.is_empty() {
                                            let exited_procs = exited_procs(&procs, &child_refs);
                                            if exited_procs.is_empty() {
                                                println!(
                                                    "[GLOBAL] All processes restarted successfully!"
                                                );
                                            } else {
                                                println!("[GLOBAL] Procs exited! {exited_procs:?}");
                                            }
                                        }
                                    }
                                    Err(e) => {
                                        let error_msg =
                                            format!("Failed to restart {}: {e}", proc.name);
                                        eprintln!("{error_msg}");
                                        failure_tx.send(error_msg).ok();
                                        starting.remove(&proc.name);
                                    }
                                }
                            });
                        }
                    }
                    Ok(Err(e)) => {
                        eprintln!("[GLOBAL] Watch error: {e}");
                    }
                    Err(_) => {
                        eprintln!("[GLOBAL] File watcher channel disconnected");
                        break;
                    }
                }
            }
        });

        Ok(())
    }

    fn shutdown(&self) {
        static SHUTDOWN_CALLED: AtomicBool = AtomicBool::new(false);

        if SHUTDOWN_CALLED.swap(true, Ordering::SeqCst) {
            println!("Shutdown already called!");
            return;
        }

        self.running.store(false, Ordering::SeqCst);

        // Stop individual processes
        for proc in self.procs.iter().rev() {
            if let Some((_, mut child)) = self.child_refs.remove(&proc.name) {
                println!("Shutting down process: {}", proc.name);
                child.kill().ok();
                child.wait().ok();
            }
            proc.cleanup();
            self.starting.remove(&proc.name);
            println!("{} stopped and cleaned up.", proc.name);
        }

        // Kill process group as fallback
        let gid = self
            .process_group_id
            .lock()
            .expect("failed to lock process group ID");
        if *gid > 0 {
            println!("Killing process group {}", *gid);
            Command::new("kill")
                .args(["-KILL", &format!("-{}", *gid)])
                .output()
                .ok();
        }

        println!("Prox shutting down.");
    }

    /// Load configuration from TOML file
    pub fn load_toml<P: AsRef<Path>>(path: P) -> anyhow::Result<Prox> {
        let content = std::fs::read_to_string(path.as_ref())
            .with_context(|| format!("Failed to read TOML file: {}", path.as_ref().display()))?;

        let prox: Prox = toml::from_str(&content)
            .with_context(|| format!("Failed to parse TOML config: {}", path.as_ref().display()))?;

        Ok(prox)
    }

    /// Load configuration from YAML file
    pub fn load_yaml<P: AsRef<Path>>(path: P) -> anyhow::Result<Prox> {
        let content = std::fs::read_to_string(path.as_ref())
            .with_context(|| format!("Failed to read YAML file: {}", path.as_ref().display()))?;

        let prox: Prox = serde_yml::from_str(&content)
            .with_context(|| format!("Failed to parse YAML config: {}", path.as_ref().display()))?;

        Ok(prox)
    }

    /// Load configuration from JSON file
    pub fn load_json<P: AsRef<Path>>(path: P) -> anyhow::Result<Prox> {
        let content = std::fs::read_to_string(path.as_ref())
            .with_context(|| format!("Failed to read JSON file: {}", path.as_ref().display()))?;

        let prox: Prox = serde_json::from_str(&content)
            .with_context(|| format!("Failed to parse JSON config: {}", path.as_ref().display()))?;

        Ok(prox)
    }
}

fn exited_procs(procs: &[Proc], child_refs: &Arc<DashMap<String, Child>>) -> Vec<String> {
    let mut exited = vec![];
    for proc in procs {
        let Some(mut child_entry) = child_refs.get_mut(&proc.name) else {
            exited.push(proc.name.clone());
            continue;
        };
        match child_entry.try_wait() {
            Ok(Some(_)) => {
                exited.push(child_entry.key().clone());
            }
            Ok(None) => {
                // Still running
            }
            Err(_) => {
                exited.push(child_entry.key().clone());
            }
        }
    }
    exited
}