use clap::Parser;
use colored::Colorize as _;
use core_foundation::base::TCFType as _;
use core_foundation::runloop::{kCFRunLoopCommonModes, CFRunLoop};
use core_foundation::string::{CFString, CFStringRef};
use core_foundation::{array::CFArray, base::CFTypeRef};
use core_graphics::event::{
    CGEventFlags, CGEventTap, CGEventTapLocation, CGEventTapOptions, CGEventTapPlacement,
    CGEventType, CallbackResult, EventField,
};
use parking_lot::Mutex;
use serde::Deserialize;
use signal_hook::consts::SIGHUP;
use signal_hook::iterator::Signals;
use std::collections::{BTreeMap, BTreeSet};
use std::io::{stdout, Write as _};
use std::path::Path;
use std::sync::{Arc, OnceLock};
use std::{ptr, thread};
use strum::{Display, EnumIter, IntoEnumIterator as _};
use termios::{tcsetattr, Termios};
use tracing::{debug, error, info, trace};

const VERSION: &str = env!("CARGO_PKG_VERSION");
const AUTHOR_URL: &str = env!("CARGO_PKG_REPOSITORY");

const LAUNCH_AGENT_PLIST: &str = include_str!("../config/ai.bma.keyswitch.plist");
const DEFAULT_CONFIG: &str = include_str!("../config/keyswitch.toml");
const LAUNCH_AGENT_PLIST_PATH: &str = "~/Library/LaunchAgents/ai.bma.keyswitch.plist";

#[derive(thiserror::Error, Debug)]
enum Error {
    #[error("Unknown input source")]
    InputSourceUnknown,
    #[error("Failed to set input source")]
    InputSourceSetFailed,
    #[error("Permission denied")]
    PermissionDenied,
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Config error: {0}")]
    Config(#[from] toml::de::Error),
}

type Result<T> = std::result::Result<T, Error>;

#[link(name = "Carbon", kind = "framework")]
unsafe extern "C" {
    fn TISCreateInputSourceList(properties: CFTypeRef, includeAll: bool) -> CFArray;
    fn TISGetInputSourceProperty(source: CFTypeRef, property_key: CFStringRef) -> CFTypeRef;
    fn TISSelectInputSource(source: CFTypeRef) -> i32;

    static kTISPropertyInputSourceID: CFStringRef;
    static kTISPropertyLocalizedName: CFStringRef;
}

struct InputSource {
    ptr_map: BTreeMap<String, (String, usize)>,
}

impl InputSource {
    fn new() -> Self {
        let sources = unsafe { TISCreateInputSourceList(ptr::null(), true) };
        let mut ptr_map = BTreeMap::new();
        for (i, src) in sources.get_all_values().iter().enumerate() {
            let id_ref = unsafe { TISGetInputSourceProperty(*src, kTISPropertyInputSourceID) };
            if id_ref.is_null() {
                continue;
            }
            let name_ref = unsafe { TISGetInputSourceProperty(*src, kTISPropertyLocalizedName) };
            if name_ref.is_null() {
                continue;
            }
            let Some(id) = cfstring_to_rust(id_ref as CFStringRef) else {
                continue;
            };
            let Some(name) = cfstring_to_rust(name_ref as CFStringRef) else {
                continue;
            };
            ptr_map.insert(id, (name, i));
        }
        InputSource { ptr_map }
    }
    fn set(&self, id: &str) -> Result<()> {
        let sources = unsafe { TISCreateInputSourceList(ptr::null(), true) };
        let Some((_, i)) = self.ptr_map.get(id) else {
            return Err(Error::InputSourceUnknown);
        };
        let src = sources
            .get((*i).try_into().unwrap())
            .ok_or(Error::InputSourceUnknown)?;
        let res = unsafe { TISSelectInputSource(*src) };
        if res != 0 {
            return Err(Error::InputSourceSetFailed);
        }
        Ok(())
    }
    fn set_logged(&self, id: &str) {
        if let Err(e) = self.set(id) {
            error!(id, error = %e);
        } else {
            debug!(id, "Input source set");
        }
    }
    fn print(&self) {
        use prettytable::{Cell, Row, Table};
        let mut table = Table::new();
        // do not make spaces between rows
        table.set_format(*prettytable::format::consts::FORMAT_NO_LINESEP);
        table.add_row(prettytable::row!["NAME", "ID"]);
        for (id, (name, _)) in &self.ptr_map {
            table.add_row(Row::new(vec![Cell::new(name), Cell::new(id)]));
        }
        table.printstd();
    }
}

fn cfstring_to_rust(cf_str: CFStringRef) -> Option<String> {
    if cf_str.is_null() {
        return None;
    }
    let s = unsafe { CFString::wrap_under_get_rule(cf_str) };
    Some(s.to_string())
}

fn disable_echo() {
    let Ok(mut termios) = Termios::from_fd(0) else {
        return;
    };
    termios.c_lflag &= !termios::ECHO;
    if let Err(e) = tcsetattr(0, termios::TCSANOW, &termios) {
        eprintln!("Failed to disable echo: {}", e);
    }
}

fn sniff_scancode() -> Result<()> {
    let current = CFRunLoop::get_current();
    let event_mask = CGEventType::KeyDown;

    disable_echo();

    println!("Sniffing scancodes... Press Ctrl+C to exit.");

    let tap = CGEventTap::new(
        CGEventTapLocation::HID,
        CGEventTapPlacement::HeadInsertEventTap,
        CGEventTapOptions::ListenOnly,
        vec![event_mask],
        move |_proxy, _tp, event| {
            let keycode = event.get_integer_value_field(EventField::KEYBOARD_EVENT_KEYCODE);
            let flags = event.get_flags();
            let mut modifiers = BTreeSet::new();
            for modifier in Modifier::iter() {
                if flags.contains(modifier.as_cg_event_flag()) {
                    modifiers.insert(modifier);
                }
            }
            println!(
                "keycode = {}, modifiers = [{}]",
                keycode,
                modifiers
                    .iter()
                    .map(|m| format!("\"{m}\""))
                    .collect::<Vec<_>>()
                    .join(", ")
            );
            stdout().flush().unwrap();
            CallbackResult::Keep
        },
    )
    .map_err(|()| Error::PermissionDenied)?;
    unsafe {
        let loop_source = tap.mach_port().create_runloop_source(0).unwrap();
        current.add_source(&loop_source, kCFRunLoopCommonModes);
        tap.enable();
        CFRunLoop::run_current();
    }
    Ok(())
}

#[derive(Parser)]
struct Args {
    #[clap(subcommand)]
    command: Option<Command>,
    #[clap(
        short,
        long,
        default_value = "~/.config/keyswitch/config.toml",
        help = "Path to config file"
    )]
    config: String,
}

#[derive(Parser)]
enum Command {
    #[clap(about = "List all available input sources")]
    ListLayouts,
    #[clap(about = "Set input source by ID")]
    Set { id: String },
    #[clap(about = "Sniff keyboard scancodes")]
    Sniff,
    #[clap(about = "Perform initial configuration")]
    Setup,
    #[clap(about = "Open config file in editor defined by $EDITOR")]
    EditConfig,
    #[clap(about = "Show program version")]
    Version,
}

#[derive(Deserialize, Debug)]
#[serde(deny_unknown_fields)]
struct Config {
    layout: LayoutConfig,
    #[serde(default)]
    rule: Vec<Rule>,
}

fn default_layout_abc() -> String {
    "com.apple.keylayout.ABC".to_string()
}

#[derive(Deserialize, Debug)]
#[serde(deny_unknown_fields)]
struct LayoutConfig {
    #[serde(default = "default_layout_abc")]
    primary: String,
    #[serde(default)]
    secondaries: Vec<String>,
}

#[derive(Deserialize, Debug)]
#[serde(deny_unknown_fields)]
struct Rule {
    keycode: i64,
    #[serde(default)]
    modifiers: BTreeSet<Modifier>,
    #[serde(skip)]
    skip_modifiers: BTreeSet<Modifier>,
    action: Action,
}

impl Rule {
    fn init(&mut self) {
        for modifier in Modifier::iter() {
            if !self.modifiers.contains(&modifier) {
                self.skip_modifiers.insert(modifier);
            }
        }
    }
    fn modifiers_match_flags(&self, flags: CGEventFlags) -> bool {
        for modifier in &self.modifiers {
            if !flags.contains(modifier.as_cg_event_flag()) {
                return false;
            }
        }
        for modifier in &self.skip_modifiers {
            if flags.contains(modifier.as_cg_event_flag()) {
                return false;
            }
        }
        true
    }
}

#[derive(Deserialize, Copy, Clone, Debug, Ord, Eq, PartialEq, PartialOrd, EnumIter, Display)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
enum Modifier {
    Shift,
    Control,
    Option,
    Command,
    NumericPad,
}

impl Modifier {
    fn as_cg_event_flag(self) -> CGEventFlags {
        match self {
            Modifier::Shift => CGEventFlags::CGEventFlagShift,
            Modifier::Control => CGEventFlags::CGEventFlagControl,
            Modifier::Option => CGEventFlags::CGEventFlagAlternate,
            Modifier::Command => CGEventFlags::CGEventFlagCommand,
            Modifier::NumericPad => CGEventFlags::CGEventFlagNumericPad,
        }
    }
}

#[derive(Deserialize, Debug)]
#[serde(tag = "kind", rename_all = "snake_case")]
#[serde(deny_unknown_fields)]
enum Action {
    Set { id: String },
    SwitchPrimarySecondaries,
    SwitchSecondaryNext,
    SwitchSecondaryPrevious,
    SwitchNext,
    SwitchPrevious,
    Command { run: String },
}

struct Context {
    config: Arc<Config>,
    input_source: Arc<InputSource>,
    primary_active: bool,
    secondary_id: usize,
}

impl Context {
    fn current_layout(&self) -> &str {
        if self.primary_active {
            &self.config.layout.primary
        } else {
            &self.config.layout.secondaries[self.secondary_id]
        }
    }
    fn set_layout(&mut self, id: &str) {
        if id == self.config.layout.primary {
            self.primary_active = true;
        } else if let Some(pos) = self
            .config
            .layout
            .secondaries
            .iter()
            .position(|secondary| secondary == id)
        {
            self.primary_active = false;
            self.secondary_id = pos;
        } else {
            eprintln!("Unknown layout ID: {}", id);
        }
    }
    fn switch_primary_secondaries(&mut self) -> &str {
        if self.config.layout.secondaries.is_empty() {
            return &self.config.layout.primary;
        }
        self.primary_active = !self.primary_active;
        self.current_layout()
    }
    fn switch_secondary_next(&mut self) {
        if self.config.layout.secondaries.is_empty() {
            return;
        }
        self.primary_active = false;
        self.secondary_id = (self.secondary_id + 1) % self.config.layout.secondaries.len();
    }
    fn switch_secondary_previous(&mut self) {
        if self.config.layout.secondaries.is_empty() {
            return;
        }
        self.primary_active = false;
        if self.secondary_id == 0 {
            self.secondary_id = self.config.layout.secondaries.len() - 1;
        } else {
            self.secondary_id -= 1;
        }
    }
    fn switch_next(&mut self) -> &str {
        if self.config.layout.secondaries.is_empty() {
            return &self.config.layout.primary;
        }
        if self.primary_active {
            self.primary_active = false;
            self.secondary_id = 0;
        } else {
            let next_id = self.secondary_id + 1;
            if next_id >= self.config.layout.secondaries.len() {
                self.primary_active = true;
                self.secondary_id = 0;
            } else {
                self.secondary_id = next_id;
            }
        }
        self.current_layout()
    }
    fn switch_previous(&mut self) -> &str {
        if self.config.layout.secondaries.is_empty() {
            return &self.config.layout.primary;
        }
        if self.primary_active {
            self.primary_active = false;
            self.secondary_id = self.config.layout.secondaries.len() - 1;
        } else if self.secondary_id == 0 {
            self.primary_active = true;
            self.secondary_id = 0;
        } else {
            self.secondary_id -= 1;
        }
        self.current_layout()
    }
}

static CONTEXT: OnceLock<Mutex<Context>> = OnceLock::new();

fn reload_context(path: &Path) -> Result<()> {
    info!("Using config file: {}", path.display());
    let config_str = std::fs::read_to_string(path)?;
    let mut config: Config = toml::from_str(&config_str)?;
    debug!(primary_layout = %config.layout.primary);
    for secondary in &config.layout.secondaries {
        debug!(secondary_layout = %secondary);
    }
    for rule in &mut config.rule {
        rule.init();
        debug!(rule = ?rule);
    }
    let context = Context {
        config: Arc::new(config),
        input_source: Arc::new(InputSource::new()),
        primary_active: true,
        secondary_id: 0,
    };
    if let Some(ctx) = CONTEXT.get() {
        let mut ctx = ctx.lock();
        *ctx = context;
    } else {
        CONTEXT
            .set(Mutex::new(context))
            .map_err(|_| "never happens")
            .unwrap();
    }
    Ok(())
}

fn serve(config_path: &Path) -> Result<()> {
    tracing_subscriber::fmt::init();
    reload_context(config_path)?;
    let mut signals = Signals::new([SIGHUP])?;
    let config_path = config_path.to_owned();
    thread::spawn(move || {
        for sig in signals.forever() {
            if sig == SIGHUP {
                info!("Received SIGHUP, reloading config");
                if let Err(e) = reload_context(&config_path) {
                    error!(error = %e, "Failed to reload config");
                } else {
                    info!("Config reloaded");
                }
            }
        }
    });
    let current = CFRunLoop::get_current();
    let tap = CGEventTap::new(
        CGEventTapLocation::HID,
        CGEventTapPlacement::HeadInsertEventTap,
        CGEventTapOptions::ListenOnly,
        vec![CGEventType::KeyDown],
        move |_proxy, _tp, event| {
            let mut context = CONTEXT.get().unwrap().lock();
            let config = context.config.clone();
            let input_source = context.input_source.clone();
            let keycode = event.get_integer_value_field(EventField::KEYBOARD_EVENT_KEYCODE);
            trace!(keycode);
            let flags = event.get_flags();
            trace!(?flags);
            for rule in &config.rule {
                trace!(rule = ?rule);
                if rule.keycode != keycode {
                    trace!("keycode does not match, skipping rule");
                    continue;
                }
                if !rule.modifiers_match_flags(flags) {
                    trace!("modifiers do not match, skipping rule");
                    continue;
                }
                trace!(curent_layout = ?context.current_layout(),
                    action = ?rule.action , "rule matched");
                match &rule.action {
                    Action::Set { id } => {
                        context.set_layout(id);
                        input_source.set_logged(id);
                    }
                    Action::SwitchPrimarySecondaries => {
                        let id = context.switch_primary_secondaries();
                        input_source.set_logged(id);
                    }
                    Action::SwitchSecondaryNext => {
                        context.switch_secondary_next();
                        let id = context.current_layout();
                        input_source.set_logged(id);
                    }
                    Action::SwitchSecondaryPrevious => {
                        context.switch_secondary_previous();
                        let id = context.current_layout();
                        input_source.set_logged(id);
                    }
                    Action::SwitchNext => {
                        let id = context.switch_next();
                        input_source.set_logged(id);
                    }
                    Action::SwitchPrevious => {
                        let id = context.switch_previous();
                        input_source.set_logged(id);
                    }
                    Action::Command { run } => {
                        debug!(command = %run, "Running command");
                        thread::spawn({
                            let run = run.clone();
                            move || {
                                if let Err(e) = std::process::Command::new("sh")
                                    .arg("-c")
                                    .arg(&run)
                                    .status()
                                {
                                    error!(command = %run, error = %e, "Failed to run command");
                                }
                            }
                        });
                    }
                }
                break;
            }
            CallbackResult::Keep
        },
    )
    .map_err(|()| Error::PermissionDenied)?;
    unsafe {
        let loop_source = tap.mach_port().create_runloop_source(0).unwrap();
        current.add_source(&loop_source, kCFRunLoopCommonModes);
        tap.enable();
        disable_echo();
        CFRunLoop::run_current();
    }
    Ok(())
}

fn write_default_config(path: &Path) -> Result<()> {
    if path.exists() {
        return Ok(());
    }
    std::fs::create_dir_all(path.parent().unwrap_or_else(|| Path::new(".")))?;
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    std::fs::write(path, DEFAULT_CONFIG)?;
    println!(
        "The default config has been written to {}",
        format!("{}", path.display()).green()
    );
    Ok(())
}

fn setup(config_path: &str) -> Result<()> {
    let current_exe = std::env::current_exe()?;
    let path = Path::new(&shellexpand::tilde(config_path).to_string()).to_owned();
    write_default_config(&path)?;
    let plist_path = Path::new(&shellexpand::tilde(LAUNCH_AGENT_PLIST_PATH).to_string()).to_owned();
    std::fs::write(
        &plist_path,
        LAUNCH_AGENT_PLIST.replace("PATH_TO_BINARY", &current_exe.display().to_string()),
    )?;
    println!(
        "The launch agent plist has been written to {}",
        format!("{}", plist_path.display()).green()
    );
    println!();
    println!(
        "Now to go Settings -> Security & Privacy -> Privacy -> Input Monitoring\nand add {} to the list.",
        format!("{}", current_exe.display()).yellow()
    );
    println!();
    println!(
        "List available keyboard layouts:\n {}",
        format!("{} list-layouts", current_exe.display()).yellow()
    );
    println!(
        "(to use the layout MUST be enabled in System Preferences -> Keyboard -> Input Sources, use ID)"
    );
    println!();
    println!(
        "Sniff key scancodes (the terminal MUST also have input monitoring allowed):\n {}",
        format!("{} sniff", current_exe.display()).yellow()
    );
    println!();
    println!(
        "Edit configuration:\n {}",
        format!("{} edit-config", current_exe.display()).yellow()
    );
    println!();
    println!(
        "Start the program:\n {}",
        format!("launchctl load {}", plist_path.display()).yellow()
    );
    println!();
    println!(
        "Reload configuration when running:\n {}",
        format!(
            "killall -1 {}",
            current_exe.file_name().unwrap().to_string_lossy()
        )
        .yellow()
    );
    println!();
    Ok(())
}

fn main() -> Result<()> {
    let args = Args::parse();
    let Some(command) = args.command else {
        let path = Path::new(&shellexpand::tilde(&args.config).to_string()).to_owned();
        return serve(&path);
    };
    match command {
        Command::ListLayouts => {
            let srcs = InputSource::new();
            srcs.print();
        }
        Command::Set { id } => {
            let srcs = InputSource::new();
            srcs.set(&id)?;
            println!("Input source set to {}", id);
        }
        Command::Sniff => {
            sniff_scancode()?;
        }
        Command::Version => {
            println!("keyswitch {} {}", VERSION, AUTHOR_URL);
        }
        Command::EditConfig => {
            let editor = std::env::var("EDITOR").unwrap_or_else(|_| "vi".to_string());
            let path = Path::new(&shellexpand::tilde(&args.config).to_string()).to_owned();
            write_default_config(&path)?;
            std::process::Command::new(editor)
                .arg(&path)
                .status()
                .map_err(Error::Io)?;
        }
        Command::Setup => {
            setup(&args.config)?;
        }
    }
    Ok(())
}
