use crate::ipc::{IpcControl, IpcNode, IpcResult, IpcWorkUnit};
use anyhow::{Context, Result};
use parking_lot::Mutex;
use std::{
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    thread,
};
use tracing::{debug, error, info};

type WorkHandler = Box<dyn Fn(&[u8]) -> Vec<u8> + Send + Sync>;
type ShutdownHandler = Box<dyn Fn() + Send + Sync>;

pub struct Program {
    handler: Option<WorkHandler>,
    shutdown_handler: Option<ShutdownHandler>,
}

impl Program {
    pub fn new() -> Self {
        Self {
            handler: None,
            shutdown_handler: None,
        }
    }

    pub fn handler<F>(mut self, f: F) -> Self
    where
        F: Fn(&[u8]) -> Vec<u8> + Send + Sync + 'static,
    {
        self.handler = Some(Box::new(f));
        self
    }

    pub fn on_shutdown<F>(mut self, f: F) -> Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        self.shutdown_handler = Some(Box::new(f));
        self
    }

    pub fn run(self) -> ! {
        if let Err(e) = self.run_inner() {
            error!("Program error: {:#}", e);
            std::process::exit(1);
        }
        std::process::exit(0);
    }

    fn run_inner(self) -> Result<()> {
        // Initialize tracing
        tracing_subscriber::fmt()
            .with_env_filter(
                tracing_subscriber::EnvFilter::try_from_default_env()
                    .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
            )
            .init();

        info!("Starting Eridani HPC program");

        let handler = self
            .handler
            .context("No work handler provided - call .handler() before .run()")?;

        // Parse args to get IPC name
        let args: Vec<String> = std::env::args().collect();
        let ipc_name = args
            .iter()
            .position(|a| a == "--ipc-name")
            .and_then(|i| args.get(i + 1))
            .context("Missing --ipc-name argument")?
            .clone();

        info!("Connecting to agent IPC: {}", ipc_name);

        // Create IPC node
        let node = IpcNode::new().context("Failed to create IPC node")?;

        // Work channel: Agent -> Program (subscribe to work)
        let work_subscriber = node
            .create_subscriber(&format!("eridani.{}.work", ipc_name))
            .context("Failed to create work subscriber")?;

        // Results channel: Program -> Agent (publish results)
        let results_publisher = node
            .create_publisher(&format!("eridani.{}.results", ipc_name))
            .context("Failed to create results publisher")?;

        // Control channel: bidirectional
        let control_pub = node
            .create_publisher(&format!("eridani.{}.control", ipc_name))
            .context("Failed to create control publisher")?;
        let control_sub = node
            .create_subscriber(&format!("eridani.{}.control", ipc_name))
            .context("Failed to create control subscriber")?;

        // Send health check response
        info!("Sending health check response");
        let health_msg = IpcControl::HealthCheckResponse { ok: true };
        control_pub
            .send_serialized(&health_msg)
            .context("Failed to send health check")?;

        // Wait for core configuration
        info!("Waiting for core configuration from agent");
        let num_cores = loop {
            if let Some(msg) = control_sub
                .try_receive_serialized::<IpcControl>()
                .context("Failed to receive control message")?
            {
                match msg {
                    IpcControl::ConfigureCores { cores } => {
                        info!("Agent configured {} worker cores", cores);
                        break cores;
                    }
                    IpcControl::Shutdown => {
                        info!("Received shutdown before configuration");
                        if let Some(handler) = &self.shutdown_handler {
                            handler();
                        }
                        return Ok(());
                    }
                    _ => {}
                }
            }
            thread::sleep(std::time::Duration::from_millis(10));
        };

        // Set up worker threads
        let shutdown = Arc::new(AtomicBool::new(false));
        let work_queue: Arc<Mutex<Vec<IpcWorkUnit>>> = Arc::new(Mutex::new(Vec::new()));
        let handler = Arc::new(handler);
        let results_publisher = Arc::new(results_publisher);
        let mut worker_threads = Vec::new();

        info!("Spawning {} worker threads", num_cores);

        // Spawn worker threads
        for core_id in 0..num_cores {
            let handler = handler.clone();
            let work_queue = work_queue.clone();
            let shutdown = shutdown.clone();
            let results_publisher = results_publisher.clone();

            let thread = thread::spawn(move || {
                // Pin thread to core (core_id + 1 to skip core 0)
                let core_ids = core_affinity::get_core_ids().unwrap_or_default();
                if let Some(core) = core_ids.get((core_id + 1) as usize) {
                    core_affinity::set_for_current(*core);
                    debug!("Worker {} pinned to core {}", core_id, core_id + 1);
                }

                // Worker loop
                loop {
                    if shutdown.load(Ordering::Relaxed) {
                        break;
                    }

                    // Get work unit from queue
                    let work_unit = {
                        let mut queue = work_queue.lock();
                        queue.pop()
                    };

                    if let Some(work_unit) = work_unit {
                        // Process work unit
                        info!(
                            worker = core_id,
                            work_unit = %work_unit.work_unit_id,
                            job = %work_unit.job_id,
                            data_bytes = work_unit.data.len(),
                            "Processing work unit"
                        );

                        let result_data = handler(&work_unit.data);

                        // Send result
                        let result = IpcResult {
                            work_unit_id: work_unit.work_unit_id.clone(),
                            job_id: work_unit.job_id.clone(),
                            data: result_data,
                        };

                        info!(
                            worker = core_id,
                            work_unit = %work_unit.work_unit_id,
                            result_bytes = result.data.len(),
                            "Work unit completed"
                        );

                        if let Err(e) = results_publisher.send_serialized(&result) {
                            error!("Failed to send result: {}", e);
                        }
                    } else {
                        // No work, sleep briefly
                        thread::sleep(std::time::Duration::from_millis(1));
                    }
                }
            });

            worker_threads.push(thread);
        }

        info!("Ready for work units");

        // Main loop: receive work units and distribute to workers
        loop {
            // Check for control messages
            while let Some(msg) = control_sub.try_receive_serialized::<IpcControl>()? {
                match msg {
                    IpcControl::Shutdown => {
                        info!("Received shutdown signal");
                        shutdown.store(true, Ordering::Relaxed);

                        // Wait for workers to finish
                        for thread in worker_threads {
                            let _ = thread.join();
                        }

                        if let Some(handler) = &self.shutdown_handler {
                            handler();
                        }

                        return Ok(());
                    }
                    IpcControl::HealthCheck => {
                        let health_msg = IpcControl::HealthCheckResponse { ok: true };
                        let _ = control_pub.send_serialized(&health_msg);
                    }
                    _ => {}
                }
            }

            // Receive work units
            while let Some(work_unit) = work_subscriber.try_receive_serialized::<IpcWorkUnit>()? {
                debug!("Received work unit {}", work_unit.work_unit_id);
                work_queue.lock().push(work_unit);
            }

            thread::sleep(std::time::Duration::from_millis(1));
        }
    }
}

impl Default for Program {
    fn default() -> Self {
        Self::new()
    }
}
