// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use std::ffi::{OsString, c_char, c_void};
use std::fs::{self, File};
use std::mem::MaybeUninit;
use std::os::windows::io::{AsRawHandle as _, FromRawHandle};
use std::path::{Path, PathBuf};
use std::ptr::{self, NonNull, null, null_mut};
use std::{io, mem, time};

use stdext::arena::{Arena, scratch_arena};
use stdext::arena_write_fmt;
use stdext::collections::{BString, BVec};
use windows_sys::Win32::Storage::FileSystem;
use windows_sys::Win32::System::{Console, IO, LibraryLoader, Threading};
use windows_sys::Win32::{Foundation, Globalization};
use windows_sys::core::*;

use crate::helpers::*;

macro_rules! w_env {
    ($s:literal) => {{
        const INPUT: &[u8] = env!($s).as_bytes();
        const OUTPUT_LEN: usize = windows_sys::core::utf16_len(INPUT) + 1;
        const OUTPUT: &[u16; OUTPUT_LEN] = {
            let mut buffer = [0; OUTPUT_LEN];
            let mut input_pos = 0;
            let mut output_pos = 0;
            while let Some((mut code_point, new_pos)) =
                windows_sys::core::decode_utf8_char(INPUT, input_pos)
            {
                input_pos = new_pos;
                if code_point <= 0xffff {
                    buffer[output_pos] = code_point as u16;
                    output_pos += 1;
                } else {
                    code_point -= 0x10000;
                    buffer[output_pos] = 0xd800 + (code_point >> 10) as u16;
                    output_pos += 1;
                    buffer[output_pos] = 0xdc00 + (code_point & 0x3ff) as u16;
                    output_pos += 1;
                }
            }
            &{ buffer }
        };
        OUTPUT.as_ptr()
    }};
}

type ReadConsoleInputExW = unsafe extern "system" fn(
    h_console_input: Foundation::HANDLE,
    lp_buffer: *mut Console::INPUT_RECORD,
    n_length: u32,
    lp_number_of_events_read: *mut u32,
    w_flags: u16,
) -> BOOL;

unsafe extern "system" fn read_console_input_ex_placeholder(
    _: Foundation::HANDLE,
    _: *mut Console::INPUT_RECORD,
    _: u32,
    _: *mut u32,
    _: u16,
) -> BOOL {
    panic!();
}

const CONSOLE_READ_NOWAIT: u16 = 0x0002;
const INVALID_CONSOLE_MODE: u32 = u32::MAX;

struct State {
    read_console_input_ex: ReadConsoleInputExW,
    stdin: Foundation::HANDLE,
    stdout: Foundation::HANDLE,
    stdin_cp_old: u32,
    stdout_cp_old: u32,
    stdin_mode_old: u32,
    stdout_mode_old: u32,
    leading_surrogate: u16,
    inject_resize: bool,
    wants_exit: bool,
}

static mut STATE: State = State {
    read_console_input_ex: read_console_input_ex_placeholder,
    stdin: null_mut(),
    stdout: null_mut(),
    stdin_cp_old: 0,
    stdout_cp_old: 0,
    stdin_mode_old: INVALID_CONSOLE_MODE,
    stdout_mode_old: INVALID_CONSOLE_MODE,
    leading_surrogate: 0,
    inject_resize: false,
    wants_exit: false,
};

extern "system" fn console_ctrl_handler(_ctrl_type: u32) -> BOOL {
    unsafe {
        STATE.wants_exit = true;
        IO::CancelIoEx(STATE.stdin, null());
    }
    1
}

/// Initializes the platform-specific state.
pub fn init() -> Deinit {
    unsafe {
        // Get the stdin and stdout handles first, so that if this function fails,
        // we at least got something to use for `write_stdout`.
        STATE.stdin = Console::GetStdHandle(Console::STD_INPUT_HANDLE);
        STATE.stdout = Console::GetStdHandle(Console::STD_OUTPUT_HANDLE);

        Deinit
    }
}

/// Reopen stdin if it's redirected (= piped input).
pub fn reopen_stdin_if_redirected() -> io::Result<Option<File>> {
    unsafe {
        let stdin = STATE.stdin;

        if stdin != Foundation::INVALID_HANDLE_VALUE
            && FileSystem::GetFileType(stdin) == FileSystem::FILE_TYPE_CHAR
        {
            return Ok(None); // stdin refers to a TTY
        }

        STATE.stdin = FileSystem::CreateFileW(
            w!("CONIN$"),
            Foundation::GENERIC_READ | Foundation::GENERIC_WRITE,
            FileSystem::FILE_SHARE_READ | FileSystem::FILE_SHARE_WRITE,
            null_mut(),
            FileSystem::OPEN_EXISTING,
            0,
            null_mut(),
        );
        if STATE.stdin == Foundation::INVALID_HANDLE_VALUE {
            return Err(last_os_error());
        }

        if stdin != Foundation::INVALID_HANDLE_VALUE {
            Ok(Some(File::from_raw_handle(stdin)))
        } else {
            Ok(None)
        }
    }
}

/// Switches the terminal into raw mode, etc.
pub fn switch_modes() -> io::Result<()> {
    unsafe {
        // `kernel32.dll` doesn't exist on OneCore variants of Windows.
        // NOTE: `kernelbase.dll` is NOT a stable API to rely on. In our case it's the best option though.
        //
        // This is written as two nested `match` statements so that we can return the error from the first
        // `load_read_func` call if it fails. The kernel32.dll lookup may contain some valid information,
        // while the kernelbase.dll lookup may not, since it's not a stable API.
        unsafe fn load_read_func(module: *const u16) -> io::Result<ReadConsoleInputExW> {
            unsafe {
                get_module(module)
                    .and_then(|m| get_proc_address(m, c"ReadConsoleInputExW".as_ptr()))
            }
        }
        STATE.read_console_input_ex = match load_read_func(w!("kernel32.dll")) {
            Ok(func) => func,
            Err(err) => match load_read_func(w!("kernelbase.dll")) {
                Ok(func) => func,
                Err(_) => return Err(err),
            },
        };

        if ptr::eq(STATE.stdin, Foundation::INVALID_HANDLE_VALUE)
            || ptr::eq(STATE.stdout, Foundation::INVALID_HANDLE_VALUE)
        {
            return Err(last_os_error());
        }

        check_bool_return(Console::GetConsoleMode(STATE.stdin, &raw mut STATE.stdin_mode_old))?;
        check_bool_return(Console::GetConsoleMode(STATE.stdout, &raw mut STATE.stdout_mode_old))?;

        match check_bool_return(Console::SetConsoleMode(
            STATE.stdin,
            Console::ENABLE_WINDOW_INPUT
                | Console::ENABLE_EXTENDED_FLAGS
                | Console::ENABLE_VIRTUAL_TERMINAL_INPUT,
        )) {
            Err(e) if e.kind() == io::ErrorKind::InvalidInput => {
                Err(io::Error::other("This application does not support the legacy console."))
            }
            other => other,
        }?;
        check_bool_return(Console::SetConsoleMode(
            STATE.stdout,
            Console::ENABLE_PROCESSED_OUTPUT
                | Console::ENABLE_WRAP_AT_EOL_OUTPUT
                | Console::ENABLE_VIRTUAL_TERMINAL_PROCESSING
                | Console::DISABLE_NEWLINE_AUTO_RETURN,
        ))?;

        check_bool_return(Console::SetConsoleCtrlHandler(Some(console_ctrl_handler), 1))?;

        STATE.stdin_cp_old = Console::GetConsoleCP();
        STATE.stdout_cp_old = Console::GetConsoleOutputCP();

        check_bool_return(Console::SetConsoleCP(Globalization::CP_UTF8))?;
        check_bool_return(Console::SetConsoleOutputCP(Globalization::CP_UTF8))?;

        Ok(())
    }
}

pub struct Deinit;

impl Drop for Deinit {
    fn drop(&mut self) {
        unsafe {
            if STATE.stdin_cp_old != 0 {
                Console::SetConsoleCP(STATE.stdin_cp_old);
                STATE.stdin_cp_old = 0;
            }
            if STATE.stdout_cp_old != 0 {
                Console::SetConsoleOutputCP(STATE.stdout_cp_old);
                STATE.stdout_cp_old = 0;
            }
            if STATE.stdin_mode_old != INVALID_CONSOLE_MODE {
                Console::SetConsoleMode(STATE.stdin, STATE.stdin_mode_old);
                STATE.stdin_mode_old = INVALID_CONSOLE_MODE;
            }
            if STATE.stdout_mode_old != INVALID_CONSOLE_MODE {
                Console::SetConsoleMode(STATE.stdout, STATE.stdout_mode_old);
                STATE.stdout_mode_old = INVALID_CONSOLE_MODE;
            }
        }
    }
}

/// During startup we need to get the window size from the terminal.
/// Because I didn't want to type a bunch of code, this function tells
/// [`read_stdin`] to inject a fake sequence, which gets picked up by
/// the input parser and provided to the TUI code.
pub fn inject_window_size_into_stdin() {
    unsafe {
        STATE.inject_resize = true;
    }
}

fn get_console_size() -> Option<Size> {
    unsafe {
        let mut info: Console::CONSOLE_SCREEN_BUFFER_INFOEX = mem::zeroed();
        info.cbSize = mem::size_of::<Console::CONSOLE_SCREEN_BUFFER_INFOEX>() as u32;
        if Console::GetConsoleScreenBufferInfoEx(STATE.stdout, &mut info) == 0 {
            return None;
        }

        let w = (info.srWindow.Right - info.srWindow.Left + 1).max(1) as CoordType;
        let h = (info.srWindow.Bottom - info.srWindow.Top + 1).max(1) as CoordType;
        Some(Size { width: w, height: h })
    }
}

/// Reads from stdin.
///
/// # Returns
///
/// * `None` if there was an error reading from stdin.
/// * `Some("")` if the given timeout was reached.
/// * Otherwise, it returns the read, non-empty string.
pub fn read_stdin(arena: &Arena, mut timeout: time::Duration) -> Option<BString<'_>> {
    let scratch = scratch_arena(Some(arena));

    // On startup we're asked to inject a window size so that the UI system can layout the elements.
    // --> Inject a fake sequence for our input parser.
    let mut resize_event = None;
    if unsafe { STATE.inject_resize } {
        unsafe { STATE.inject_resize = false };
        timeout = time::Duration::ZERO;
        resize_event = get_console_size();
    }

    let read_poll = timeout != time::Duration::MAX; // there is a timeout -> don't block in read()
    let input_buf = scratch.alloc_uninit_slice(4 * KIBI);
    let mut input_buf_cap = input_buf.len();
    let utf16_buf = scratch.alloc_uninit_slice(4 * KIBI);
    let mut utf16_buf_len = 0;

    // If there was a leftover leading surrogate from the last read, we prepend it to the buffer.
    if unsafe { STATE.leading_surrogate } != 0 {
        utf16_buf[0] = MaybeUninit::new(unsafe { STATE.leading_surrogate });
        utf16_buf_len = 1;
        input_buf_cap -= 1;
        unsafe { STATE.leading_surrogate = 0 };
    }

    // Read until there's either a timeout or we have something to process.
    loop {
        if timeout != time::Duration::MAX {
            let beg = time::Instant::now();

            match unsafe { Threading::WaitForSingleObject(STATE.stdin, timeout.as_millis() as u32) }
            {
                // Ready to read? Continue with reading below.
                Foundation::WAIT_OBJECT_0 => {}
                // Timeout? Skip reading entirely.
                Foundation::WAIT_TIMEOUT => break,
                // Error? Tell the caller stdin is broken.
                _ => return None,
            }

            timeout = timeout.saturating_sub(beg.elapsed());
        }

        // Read from stdin.
        let input = unsafe {
            // If we had a `inject_resize`, we don't want to block indefinitely for other pending input on startup,
            // but are still interested in any other pending input that may be waiting for us.
            let flags = if read_poll { CONSOLE_READ_NOWAIT } else { 0 };
            let mut read = 0;
            let ok = (STATE.read_console_input_ex)(
                STATE.stdin,
                input_buf[0].as_mut_ptr(),
                input_buf_cap as u32,
                &mut read,
                flags,
            );
            if ok == 0 || STATE.wants_exit {
                return None;
            }
            input_buf[..read as usize].assume_init_ref()
        };

        // Convert Win32 input records into UTF16.
        for inp in input {
            match inp.EventType as u32 {
                Console::KEY_EVENT => {
                    let event = unsafe { &inp.Event.KeyEvent };
                    let ch = unsafe { event.uChar.UnicodeChar };
                    if event.bKeyDown != 0 && ch != 0 {
                        utf16_buf[utf16_buf_len] = MaybeUninit::new(ch);
                        utf16_buf_len += 1;
                    }
                }
                Console::WINDOW_BUFFER_SIZE_EVENT => {
                    let event = unsafe { &inp.Event.WindowBufferSizeEvent };
                    let w = event.dwSize.X as CoordType;
                    let h = event.dwSize.Y as CoordType;
                    // Windows is prone to sending broken/useless `WINDOW_BUFFER_SIZE_EVENT`s.
                    // E.g. starting conhost will emit 3 in a row. Skip rendering in that case.
                    if w > 0 && h > 0 {
                        resize_event = Some(Size { width: w, height: h });
                    }
                }
                _ => {}
            }
        }

        if resize_event.is_some() || utf16_buf_len != 0 {
            break;
        }
    }

    const RESIZE_EVENT_FMT_MAX_LEN: usize = 16; // "\x1b[8;65535;65535t"
    let resize_event_len = if resize_event.is_some() { RESIZE_EVENT_FMT_MAX_LEN } else { 0 };
    // +1 to account for a potential `STATE.leading_surrogate`.
    let utf8_max_len = (utf16_buf_len + 1) * 3;
    let mut text = BString::empty();
    text.reserve(arena, utf8_max_len + resize_event_len);

    // Now prepend our previously extracted resize event.
    if let Some(resize_event) = resize_event {
        // If I read xterm's documentation correctly, CSI 18 t reports the window size in characters.
        // CSI 8 ; height ; width t is the response. Of course, we didn't send the request,
        // but we can use this fake response to trigger the editor to resize itself.
        arena_write_fmt!(arena, text, "\x1b[8;{};{}t", resize_event.height, resize_event.width);
    }

    // If the input ends with a lone lead surrogate, we need to remember it for the next read.
    if utf16_buf_len > 0 {
        unsafe {
            let last_char = utf16_buf[utf16_buf_len - 1].assume_init();
            if (0xD800..0xDC00).contains(&last_char) {
                STATE.leading_surrogate = last_char;
                utf16_buf_len -= 1;
            }
        }
    }

    // Convert the remaining input to UTF8, the sane encoding.
    if utf16_buf_len > 0 {
        unsafe {
            let vec = text.as_mut_vec();
            let spare = vec.spare_capacity_mut();

            let len = Globalization::WideCharToMultiByte(
                Globalization::CP_UTF8,
                0,
                utf16_buf[0].as_ptr(),
                utf16_buf_len as i32,
                spare.as_mut_ptr() as *mut _,
                spare.len() as i32,
                null(),
                null_mut(),
            );

            if len > 0 {
                vec.set_len(vec.len() + len as usize);
            }
        }
    }

    Some(text)
}

/// Writes a string to stdout.
///
/// Use this instead of `print!` or `println!` to avoid
/// the overhead of Rust's stdio handling. Don't need that.
pub fn write_stdout(text: &str) {
    unsafe {
        let mut offset = 0;

        while offset < text.len() {
            let ptr = text.as_ptr().add(offset);
            let write = (text.len() - offset).min(GIBI) as u32;
            let mut written = 0;
            let ok = FileSystem::WriteFile(STATE.stdout, ptr, write, &mut written, null_mut());
            offset += written as usize;
            if ok == 0 || written == 0 {
                break;
            }
        }
    }
}

pub fn drives() -> impl Iterator<Item = char> {
    unsafe {
        let mut mask = FileSystem::GetLogicalDrives();
        std::iter::from_fn(move || {
            let bit = mask.trailing_zeros();
            if bit >= 26 {
                None
            } else {
                mask &= !(1 << bit);
                Some((b'A' + bit as u8) as char)
            }
        })
    }
}

/// A unique identifier for a file.
pub enum FileId {
    Id(FileSystem::FILE_ID_INFO),
    Path(PathBuf),
}

impl PartialEq for FileId {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Id(left), Self::Id(right)) => {
                // Lowers to an efficient word-wise comparison.
                const SIZE: usize = std::mem::size_of::<FileSystem::FILE_ID_INFO>();
                let a: &[u8; SIZE] = unsafe { mem::transmute(left) };
                let b: &[u8; SIZE] = unsafe { mem::transmute(right) };
                a == b
            }
            (Self::Path(left), Self::Path(right)) => left == right,
            _ => false,
        }
    }
}

impl Eq for FileId {}

/// Returns a unique identifier for the given file by handle or path.
pub fn file_id(file: Option<&File>, path: &Path) -> io::Result<FileId> {
    let file = match file {
        Some(f) => f,
        None => &File::open(path)?,
    };

    file_id_from_handle(file).or_else(|_| Ok(FileId::Path(std::fs::canonicalize(path)?)))
}

fn file_id_from_handle(file: &File) -> io::Result<FileId> {
    unsafe {
        let mut info = MaybeUninit::<FileSystem::FILE_ID_INFO>::uninit();
        check_bool_return(FileSystem::GetFileInformationByHandleEx(
            file.as_raw_handle(),
            FileSystem::FileIdInfo,
            info.as_mut_ptr() as *mut _,
            mem::size_of::<FileSystem::FILE_ID_INFO>() as u32,
        ))?;
        Ok(FileId::Id(info.assume_init()))
    }
}

/// Canonicalizes the given path.
///
/// This differs from [`fs::canonicalize`] in that it strips the `\\?\` UNC
/// prefix on Windows. This is because it's confusing/ugly when displaying it.
pub fn canonicalize(path: &Path) -> std::io::Result<PathBuf> {
    let mut path = fs::canonicalize(path)?;
    let path = path.as_mut_os_string();
    let mut path = mem::take(path).into_encoded_bytes();

    if path.len() > 6 && &path[0..4] == br"\\?\" && path[4].is_ascii_uppercase() && path[5] == b':'
    {
        path.drain(0..4);
    }

    let path = unsafe { OsString::from_encoded_bytes_unchecked(path) };
    let path = PathBuf::from(path);
    Ok(path)
}

unsafe fn get_module(name: *const u16) -> io::Result<NonNull<c_void>> {
    unsafe { check_ptr_return(LibraryLoader::GetModuleHandleW(name)) }
}

unsafe fn load_library(name: *const u16) -> io::Result<NonNull<c_void>> {
    unsafe {
        check_ptr_return(LibraryLoader::LoadLibraryExW(
            name,
            null_mut(),
            LibraryLoader::LOAD_LIBRARY_SEARCH_SYSTEM32,
        ))
    }
}

/// Loads a function from a dynamic library.
///
/// # Safety
///
/// This function is highly unsafe as it requires you to know the exact type
/// of the function you're loading. No type checks whatsoever are performed.
//
// It'd be nice to constrain T to std::marker::FnPtr, but that's unstable.
pub unsafe fn get_proc_address<T>(handle: NonNull<c_void>, name: *const c_char) -> io::Result<T> {
    unsafe {
        let ptr = LibraryLoader::GetProcAddress(handle.as_ptr(), name as *const u8);
        if let Some(ptr) = ptr { Ok(mem::transmute_copy(&ptr)) } else { Err(last_os_error()) }
    }
}

pub struct LibIcu {
    pub libicuuc: NonNull<c_void>,
    pub libicui18n: NonNull<c_void>,
}

pub fn load_icu() -> io::Result<LibIcu> {
    const fn const_ptr_u16_eq(a: *const u16, b: *const u16) -> bool {
        unsafe {
            let mut a = a;
            let mut b = b;
            loop {
                if *a != *b {
                    return false;
                }
                if *a == 0 {
                    return true;
                }
                a = a.add(1);
                b = b.add(1);
            }
        }
    }

    const LIBICUUC: *const u16 = w_env!("EDIT_CFG_ICUUC_SONAME");
    const LIBICUI18N: *const u16 = w_env!("EDIT_CFG_ICUI18N_SONAME");

    if const { const_ptr_u16_eq(LIBICUUC, LIBICUI18N) } {
        let icu = unsafe { load_library(LIBICUUC)? };
        Ok(LibIcu { libicuuc: icu, libicui18n: icu })
    } else {
        let libicuuc = unsafe { load_library(LIBICUUC)? };
        let libicui18n = unsafe { load_library(LIBICUI18N)? };
        Ok(LibIcu { libicuuc, libicui18n })
    }
}

/// Returns a list of preferred languages for the current user.
pub fn preferred_languages<'a>(arena: &'a Arena) -> BVec<'a, &'a str> {
    // If the GetUserPreferredUILanguages() don't fit into 512 characters,
    // honestly, just give up. How many languages do you realistically need?
    const LEN: usize = 512;

    let scratch = scratch_arena(Some(arena));

    // Get the list of preferred languages via `GetUserPreferredUILanguages`.
    let langs = unsafe {
        let buf = scratch.alloc_uninit_slice(LEN);
        let mut len = buf.len() as u32;
        let mut num = 0;

        let ok = Globalization::GetUserPreferredUILanguages(
            Globalization::MUI_LANGUAGE_NAME,
            &mut num,
            buf[0].as_mut_ptr(),
            &mut len,
        );

        if ok == 0 || num == 0 {
            len = 0;
        }

        // Drop the terminating double-null character.
        len = len.saturating_sub(1);

        buf[..len as usize].assume_init_ref()
    };

    // Convert UTF16 to UTF8.
    let langs = BString::from_utf16_lossy(arena, langs).leak();

    // Split the null-delimited string into individual chunks
    // and copy them into the given arena.
    let mut res = BVec::empty();
    res.extend_sloppy(arena, langs.split_terminator('\0').filter(|s| !s.is_empty()));
    res
}

#[inline]
#[cold]
fn last_os_error() -> io::Error {
    io::Error::last_os_error()
}

fn check_bool_return(ret: BOOL) -> io::Result<()> {
    if ret == 0 { Err(last_os_error()) } else { Ok(()) }
}

fn check_ptr_return<T>(ret: *mut T) -> io::Result<NonNull<T>> {
    NonNull::new(ret).ok_or_else(last_os_error)
}
