use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use colored::Colorize;
use futures_util::{SinkExt, StreamExt};
use native_tls::TlsConnector;
use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;
use tokio_tungstenite::{connect_async_tls_with_config, tungstenite::Message, Connector};

use base64::Engine;
use helios_common::{ClientMessage, ServerMessage};
use uuid::Uuid;

mod shell;
mod screenshot;
mod input;
mod windows_mgmt;

// ─── CLI Output ─────────────────────────────────────────────────────────────

/// Truncate a string to at most `max_chars` Unicode characters.
fn trunc(s: &str, max_chars: usize) -> String {
    let mut chars = s.chars();
    let truncated: String = chars.by_ref().take(max_chars).collect();
    if chars.next().is_some() {
        format!("{}…", truncated)
    } else {
        truncated
    }
}

fn banner() {
    println!();
    println!("  {}  HELIOS REMOTE  v{} ({})", "☀".yellow().bold(), env!("CARGO_PKG_VERSION"), env!("GIT_COMMIT"));
    #[cfg(windows)]
    {
        let admin = is_admin();
        let (icon, admin_str) = if admin {
            ("👑", "admin".green().bold().to_string())
        } else {
            ("👤", "user (no admin)".yellow().to_string())
        };
        println!("  {}  {}", icon, admin_str);
    }
    println!();
}

#[cfg(windows)]
fn is_admin() -> bool {
    use windows::Win32::UI::Shell::IsUserAnAdmin;
    unsafe { IsUserAnAdmin().as_bool() }
}

#[cfg(windows)]
fn enable_ansi() {
    use windows::Win32::System::Console::{
        GetConsoleMode, GetStdHandle, SetConsoleMode,
        ENABLE_VIRTUAL_TERMINAL_PROCESSING, STD_OUTPUT_HANDLE,
    };
    unsafe {
        if let Ok(handle) = GetStdHandle(STD_OUTPUT_HANDLE) {
            let mut mode = Default::default();
            if GetConsoleMode(handle, &mut mode).is_ok() {
                let _ = SetConsoleMode(handle, mode | ENABLE_VIRTUAL_TERMINAL_PROCESSING);
            }
        }
    }
}

macro_rules! log_status {
    ($($arg:tt)*) => {
        println!("  {}", format!($($arg)*));
    };
}

macro_rules! log_ok {
    ($($arg:tt)*) => {
        println!("  {} {}", "✅", format!($($arg)*));
    };
}

macro_rules! log_err {
    ($($arg:tt)*) => {
        println!("  {} {}", "✗".red().bold(), format!($($arg)*));
    };
}

macro_rules! log_cmd {
    ($emoji:expr, $($arg:tt)*) => {
        println!("  {}  {}", $emoji, format!($($arg)*));
    };
}

// ────────────────────────────────────────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
struct Config {
    relay_url: String,
    api_key: String,
    label: Option<String>,
    session_id: Option<String>, // persistent UUID
}

impl Config {
    fn config_path() -> PathBuf {
        let base = dirs::config_dir()
            .or_else(|| dirs::home_dir())
            .unwrap_or_else(|| PathBuf::from("."));
        base.join("helios-remote").join("config.toml")
    }

    fn load() -> Option<Self> {
        let path = Self::config_path();
        let data = std::fs::read_to_string(&path).ok()?;
        toml::from_str(&data).ok()
    }

    fn save(&self) -> std::io::Result<()> {
        let path = Self::config_path();
        std::fs::create_dir_all(path.parent().unwrap())?;
        let data = toml::to_string_pretty(self).unwrap();
        std::fs::write(&path, data)?;
        Ok(())
    }
}

fn prompt_config() -> Config {
    use std::io::Write;

    let relay_url = {
        let default = "wss://remote.agent-helios.me/ws";
        print!("  {} Relay URL [{}]: ", "→".cyan().bold(), default);
        std::io::stdout().flush().unwrap();
        let mut input = String::new();
        std::io::stdin().read_line(&mut input).unwrap();
        let trimmed = input.trim();
        if trimmed.is_empty() {
            default.to_string()
        } else {
            trimmed.to_string()
        }
    };

    let api_key = {
        print!("  {} API Key: ", "→".cyan().bold());
        std::io::stdout().flush().unwrap();
        let mut input = String::new();
        std::io::stdin().read_line(&mut input).unwrap();
        input.trim().to_string()
    };

    let label = {
        let default_label = hostname();
        print!("  {} Label for this PC [{}]: ", "→".cyan().bold(), default_label);
        std::io::stdout().flush().unwrap();
        let mut input = String::new();
        std::io::stdin().read_line(&mut input).unwrap();
        let trimmed = input.trim().to_string();
        if trimmed.is_empty() {
            Some(default_label)
        } else {
            Some(trimmed)
        }
    };

    Config { relay_url, api_key, label, session_id: None }
}

#[tokio::main]
async fn main() {
    // Enable ANSI color codes on Windows (required when running as admin)
    #[cfg(windows)]
    enable_ansi();

    // Suppress tracing output by default
    if std::env::var("RUST_LOG").is_err() {
        unsafe { std::env::set_var("RUST_LOG", "off"); }
    }

    banner();

    // Load or prompt for config
    let config = match Config::load() {
        Some(c) => c,
        None => {
            log_status!("No config found — first-time setup");
            println!();
            let c = prompt_config();
            println!();
            if let Err(e) = c.save() {
                log_err!("Failed to save config: {e}");
            } else {
                log_ok!("Config saved");
            }
            c
        }
    };

    // Resolve or generate persistent session UUID
    let sid: Uuid = match &config.session_id {
        Some(id) => Uuid::parse_str(id).unwrap_or_else(|_| Uuid::new_v4()),
        None => {
            let id = Uuid::new_v4();
            let mut cfg = config.clone();
            cfg.session_id = Some(id.to_string());
            if let Err(e) = cfg.save() {
                log_err!("Failed to save session_id: {e}");
            }
            id
        }
    };

    let config = Arc::new(config);
    let shell = Arc::new(Mutex::new(shell::PersistentShell::new()));

    // Connect with exponential backoff
    let mut backoff = Duration::from_secs(1);
    const MAX_BACKOFF: Duration = Duration::from_secs(30);

    loop {
        let host = config.relay_url
            .trim_start_matches("wss://")
            .trim_start_matches("ws://")
            .split('/')
            .next()
            .unwrap_or(&config.relay_url);

        log_status!("🌐  Connecting to {}...", host);

        let tls_connector = TlsConnector::builder()
            .danger_accept_invalid_certs(true)
            .build()
            .expect("TLS connector build failed");
        let connector = Connector::NativeTls(tls_connector);

        match connect_async_tls_with_config(&config.relay_url, None, false, Some(connector)).await {
            Ok((ws_stream, _)) => {
                let label = config.label.clone().unwrap_or_else(|| hostname());
                log_ok!(
                    "Connected  {}  {}  {}  Session {}",
                    "·".dimmed(),
                    label.bold(),
                    "·".dimmed(),
                    sid.to_string().dimmed()
                );
                println!();
                backoff = Duration::from_secs(1);

                let (mut write, mut read) = ws_stream.split();

                // Send Hello
                let hello = ClientMessage::Hello {
                    label: config.label.clone(),
                };
                let hello_json = serde_json::to_string(&hello).unwrap();
                if let Err(e) = write.send(Message::Text(hello_json)).await {
                    log_err!("Failed to send Hello: {e}");
                    tokio::time::sleep(backoff).await;
                    backoff = (backoff * 2).min(MAX_BACKOFF);
                    continue;
                }

                let write = Arc::new(Mutex::new(write));

                while let Some(msg_result) = read.next().await {
                    match msg_result {
                        Ok(Message::Text(text)) => {
                            let server_msg: ServerMessage = match serde_json::from_str(&text) {
                                Ok(m) => m,
                                Err(e) => {
                                    log_err!("Failed to parse server message: {e}");
                                    continue;
                                }
                            };

                            let write_clone = Arc::clone(&write);
                            let shell_clone = Arc::clone(&shell);

                            tokio::spawn(async move {
                                // tokio isolates panics per task — a panic here won't kill
                                // the main loop. handle_message uses map_err everywhere so
                                // it should never panic in practice.
                                let response = handle_message(server_msg, shell_clone).await;
                                let json = match serde_json::to_string(&response) {
                                    Ok(j) => j,
                                    Err(e) => {
                                        log_err!("Failed to serialize response: {e}");
                                        return;
                                    }
                                };
                                let mut w = write_clone.lock().await;
                                if let Err(e) = w.send(Message::Text(json)).await {
                                    log_err!("Failed to send response: {e}");
                                }
                            });
                        }
                        Ok(Message::Ping(data)) => {
                            let mut w = write.lock().await;
                            let _ = w.send(Message::Pong(data)).await;
                        }
                        Ok(Message::Close(_)) => {
                            log_err!("Connection lost — reconnecting...");
                            break;
                        }
                        Err(e) => {
                            log_err!("Connection lost: {e} — reconnecting...");
                            break;
                        }
                        _ => {}
                    }
                }
            }
            Err(e) => {
                log_err!("Connection failed: {e}");
            }
        }

        tokio::time::sleep(backoff).await;
        backoff = (backoff * 2).min(MAX_BACKOFF);
    }
}

fn hostname() -> String {
    std::fs::read_to_string("/etc/hostname")
        .unwrap_or_default()
        .trim()
        .to_string()
        .or_else(|| std::env::var("COMPUTERNAME").ok())
        .unwrap_or_else(|| "unknown".to_string())
}

trait OrElseString {
    fn or_else(self, f: impl FnOnce() -> Option<String>) -> String;
    fn unwrap_or_else(self, f: impl FnOnce() -> String) -> String;
}

impl OrElseString for String {
    fn or_else(self, f: impl FnOnce() -> Option<String>) -> String {
        if self.is_empty() { f().unwrap_or_default() } else { self }
    }
    fn unwrap_or_else(self, f: impl FnOnce() -> String) -> String {
        if self.is_empty() { f() } else { self }
    }
}

async fn handle_message(
    msg: ServerMessage,
    shell: Arc<Mutex<shell::PersistentShell>>,
) -> ClientMessage {
    match msg {
        ServerMessage::WindowScreenshotRequest { request_id, window_id } => {
            log_cmd!("📷", "screenshot window {}", window_id);
            match screenshot::take_window_screenshot(window_id) {
                Ok((image_base64, width, height)) => {
                    log_ok!("Done  {}  {}×{}", "·".dimmed(), width, height);
                    ClientMessage::ScreenshotResponse { request_id, image_base64, width, height }
                }
                Err(e) => {
                    log_err!("Window screenshot failed: {e}");
                    ClientMessage::Error { request_id, message: format!("Window screenshot failed: {e}") }
                }
            }
        }

        ServerMessage::ScreenshotRequest { request_id } => {
            log_cmd!("📷", "screenshot");
            match screenshot::take_screenshot() {
                Ok((image_base64, width, height)) => {
                    log_ok!("Done  {}  {}×{}", "·".dimmed(), width, height);
                    ClientMessage::ScreenshotResponse {
                        request_id,
                        image_base64,
                        width,
                        height,
                    }
                }
                Err(e) => {
                    log_err!("Screenshot failed: {e}");
                    ClientMessage::Error {
                        request_id,
                        message: format!("Screenshot failed: {e}"),
                    }
                }
            }
        }

        ServerMessage::PromptRequest { request_id, message, title } => {
            let _title = title.unwrap_or_else(|| "Helios Remote".to_string());
            #[cfg(windows)]
            let title = _title.clone();
            log_cmd!("💬", "prompt › {}", trunc(&message, 60));
            #[cfg(windows)]
            {
                use windows::core::PCWSTR;
                use windows::Win32::UI::WindowsAndMessaging::{MessageBoxW, MB_OK, MB_ICONINFORMATION, HWND_DESKTOP};
                let msg_wide: Vec<u16> = message.encode_utf16().chain(std::iter::once(0)).collect();
                let title_wide: Vec<u16> = title.encode_utf16().chain(std::iter::once(0)).collect();
                // Run blocking MessageBox in a thread so we don't block the async runtime
                let msg_clone = message.clone();
                tokio::task::spawn_blocking(move || {
                    unsafe {
                        MessageBoxW(
                            HWND_DESKTOP,
                            PCWSTR(msg_wide.as_ptr()),
                            PCWSTR(title_wide.as_ptr()),
                            MB_OK | MB_ICONINFORMATION,
                        );
                    }
                }).await.ok();
                log_ok!("User confirmed: {}", trunc(&msg_clone, 40));
            }
            #[cfg(not(windows))]
            {
                // On non-Windows just log it
                println!("  [PROMPT] {}", message);
            }
            ClientMessage::Ack { request_id }
        }

        ServerMessage::ExecRequest { request_id, command, timeout_ms } => {
            let cmd_display = trunc(&command, 60);
            log_cmd!("⚡", "exec › {}", cmd_display);
            let mut sh = shell.lock().await;
            match sh.run(&command, timeout_ms).await {
                Ok((stdout, stderr, exit_code)) => {
                    let out = stdout.trim().lines().next().unwrap_or("").to_string();
                    let out_display = trunc(&out, 60);
                    if exit_code != 0 {
                        if out_display.is_empty() {
                            log_err!("exit {}", exit_code);
                        } else {
                            log_err!("{}  {}  exit {}", out_display, "·".dimmed(), exit_code);
                        }
                    } else if out_display.is_empty() {
                        log_ok!("exit 0");
                    } else {
                        log_ok!("{}  {}  exit 0", out_display, "·".dimmed());
                    }
                    let _ = stderr;
                    ClientMessage::ExecResponse {
                        request_id,
                        stdout,
                        stderr,
                        exit_code,
                    }
                }
                Err(e) => {
                    log_err!("exec failed: {e}");
                    ClientMessage::Error {
                        request_id,
                        message: format!(
                            "Exec failed for command {:?}.\nError: {e}",
                            command
                        ),
                    }
                }
            }
        }

        ServerMessage::ClickRequest { request_id, x, y, button } => {
            log_cmd!("🖱 ", "click ({x}, {y}) {:?}", button);
            match input::click(x, y, &button) {
                Ok(()) => {
                    log_ok!("Done");
                    ClientMessage::Ack { request_id }
                }
                Err(e) => {
                    log_err!("click failed: {e}");
                    ClientMessage::Error {
                        request_id,
                        message: format!("Click at ({x},{y}) failed: {e}"),
                    }
                }
            }
        }

        ServerMessage::TypeRequest { request_id, text } => {
            log_cmd!("⌨ ", "type {} chars", text.len());
            match input::type_text(&text) {
                Ok(()) => {
                    log_ok!("Done");
                    ClientMessage::Ack { request_id }
                }
                Err(e) => {
                    log_err!("type failed: {e}");
                    ClientMessage::Error {
                        request_id,
                        message: format!("Type failed: {e}"),
                    }
                }
            }
        }

        ServerMessage::ListWindowsRequest { request_id } => {
            log_cmd!("🪟", "list-windows");
            match windows_mgmt::list_windows() {
                Ok(windows) => {
                    log_ok!("{} windows", windows.len());
                    ClientMessage::ListWindowsResponse { request_id, windows }
                }
                Err(e) => {
                    log_err!("list-windows failed: {e}");
                    ClientMessage::Error { request_id, message: e }
                }
            }
        }

        ServerMessage::MinimizeAllRequest { request_id } => {
            log_cmd!("🪟", "minimize-all");
            match windows_mgmt::minimize_all() {
                Ok(()) => {
                    log_ok!("Done");
                    ClientMessage::Ack { request_id }
                }
                Err(e) => {
                    log_err!("minimize-all failed: {e}");
                    ClientMessage::Error { request_id, message: e }
                }
            }
        }

        ServerMessage::FocusWindowRequest { request_id, window_id } => {
            log_cmd!("🪟", "focus-window {window_id}");
            match windows_mgmt::focus_window(window_id) {
                Ok(()) => {
                    log_ok!("Done");
                    ClientMessage::Ack { request_id }
                }
                Err(e) => {
                    log_err!("focus-window failed: {e}");
                    ClientMessage::Error { request_id, message: e }
                }
            }
        }

        ServerMessage::MaximizeAndFocusRequest { request_id, window_id } => {
            log_cmd!("🪟", "maximize-and-focus {window_id}");
            match windows_mgmt::maximize_and_focus(window_id) {
                Ok(()) => {
                    log_ok!("Done");
                    ClientMessage::Ack { request_id }
                }
                Err(e) => {
                    log_err!("maximize-and-focus failed: {e}");
                    ClientMessage::Error { request_id, message: e }
                }
            }
        }

        ServerMessage::VersionRequest { request_id } => {
            log_cmd!("ℹ ", "version");
            ClientMessage::VersionResponse {
                request_id,
                version: env!("CARGO_PKG_VERSION").to_string(),
                commit: env!("GIT_COMMIT").to_string(),
            }
        }

        ServerMessage::UploadRequest { request_id, path, content_base64 } => {
            log_cmd!("⬆ ", "upload → {}", path);
            match (|| -> Result<(), String> {
                let bytes = base64::engine::general_purpose::STANDARD
                    .decode(&content_base64)
                    .map_err(|e| format!("base64 decode: {e}"))?;
                if let Some(parent) = std::path::Path::new(&path).parent() {
                    std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
                }
                std::fs::write(&path, &bytes).map_err(|e| e.to_string())?;
                Ok(())
            })() {
                Ok(()) => {
                    log_ok!("Saved {}", path);
                    ClientMessage::Ack { request_id }
                }
                Err(e) => {
                    log_err!("upload failed: {e}");
                    ClientMessage::Error { request_id, message: e }
                }
            }
        }

        ServerMessage::DownloadRequest { request_id, path } => {
            log_cmd!("⬇ ", "download ← {}", path);
            match std::fs::read(&path) {
                Ok(bytes) => {
                    let size = bytes.len() as u64;
                    let content_base64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
                    log_ok!("Sent {} bytes", size);
                    ClientMessage::DownloadResponse { request_id, content_base64, size }
                }
                Err(e) => {
                    log_err!("download failed: {e}");
                    ClientMessage::Error { request_id, message: format!("Read failed: {e}") }
                }
            }
        }

        ServerMessage::RunRequest { request_id, program, args } => {
            log_cmd!("🚀", "run › {}", program);
            use std::process::Command as StdCommand;
            match StdCommand::new(&program).args(&args).spawn() {
                Ok(_) => {
                    log_ok!("Started {}", program);
                    ClientMessage::Ack { request_id }
                }
                Err(e) => {
                    log_err!("run failed: {e}");
                    ClientMessage::Error { request_id, message: format!("Failed to start '{}': {e}", program) }
                }
            }
        }

        ServerMessage::ClipboardGetRequest { request_id } => {
            log_cmd!("📋", "clipboard-get");
            let out = tokio::process::Command::new("powershell.exe")
                .args(["-NoProfile", "-NonInteractive", "-Command", "Get-Clipboard"])
                .output().await;
            match out {
                Ok(o) => {
                    let text = String::from_utf8_lossy(&o.stdout).trim().to_string();
                    log_ok!("Got {} chars", text.len());
                    ClientMessage::ClipboardGetResponse { request_id, text }
                }
                Err(e) => ClientMessage::Error { request_id, message: format!("Clipboard get failed: {e}") }
            }
        }

        ServerMessage::ClipboardSetRequest { request_id, text } => {
            log_cmd!("📋", "clipboard-set › {} chars", text.len());
            let cmd = format!("Set-Clipboard -Value '{}'", text.replace('\'', "''"));
            let out = tokio::process::Command::new("powershell.exe")
                .args(["-NoProfile", "-NonInteractive", "-Command", &cmd])
                .output().await;
            match out {
                Ok(_) => {
                    log_ok!("Set clipboard");
                    ClientMessage::Ack { request_id }
                }
                Err(e) => ClientMessage::Error { request_id, message: format!("Clipboard set failed: {e}") }
            }
        }

        ServerMessage::Ack { request_id } => {
            ClientMessage::Ack { request_id }
        }

        ServerMessage::Error { request_id, message } => {
            log_err!("server error: {message}");
            if let Some(rid) = request_id {
                ClientMessage::Ack { request_id: rid }
            } else {
                ClientMessage::Hello { label: None }
            }
        }
    }
}