use num_derive::{FromPrimitive, ToPrimitive};
use num_traits::FromPrimitive;
use std::fmt;
use std::fmt::{Arguments, Display};
use std::io::{stderr, stdout, Write};
use std::panic::Location;
use std::process::exit;

use crate::ffi::LogLevelCxx;
use crate::{Utf8CStr, Utf8CStrBufArr};

// Error handling and logging throughout the Rust codebase in Magisk:
//
// All errors should be logged and consumed as soon as possible and converted into LoggedError.
// For `Result` with errors that implement the `Display` trait, use the `?` operator to
// log and convert to LoggedResult.
//
// To log an error with more information, use `ResultExt::log_with_msg()`.
//
// The "cxx" method variants in `CxxResultExt` are only used for C++ interop and
// should not be used directly in any Rust code.
//
// For general logging, use the <level>!(...) macros.

// Ugly hack to avoid using enum
#[allow(non_snake_case, non_upper_case_globals)]
mod LogFlag {
    pub const DisableError: u32 = 1 << 0;
    pub const DisableWarn: u32 = 1 << 1;
    pub const DisableInfo: u32 = 1 << 2;
    pub const DisableDebug: u32 = 1 << 3;
    pub const ExitOnError: u32 = 1 << 4;
}

#[derive(Copy, Clone, FromPrimitive, ToPrimitive)]
#[repr(i32)]
pub enum LogLevel {
    ErrorCxx = LogLevelCxx::ErrorCxx.repr,
    Error = LogLevelCxx::Error.repr,
    Warn = LogLevelCxx::Warn.repr,
    Info = LogLevelCxx::Info.repr,
    Debug = LogLevelCxx::Debug.repr,
}

// We don't need to care about thread safety, because all
// logger changes will only happen on the main thread.
pub static mut LOGGER: Logger = Logger {
    write: |_, _| {},
    flags: 0,
};

type LogWriter = fn(level: LogLevel, msg: &Utf8CStr);
type Formatter<'a> = &'a mut dyn fmt::Write;

#[derive(Copy, Clone)]
pub struct Logger {
    pub write: LogWriter,
    pub flags: u32,
}

pub fn exit_on_error(b: bool) {
    unsafe {
        if b {
            LOGGER.flags |= LogFlag::ExitOnError;
        } else {
            LOGGER.flags &= !LogFlag::ExitOnError;
        }
    }
}

impl LogLevel {
    fn as_disable_flag(&self) -> u32 {
        match *self {
            LogLevel::Error | LogLevel::ErrorCxx => LogFlag::DisableError,
            LogLevel::Warn => LogFlag::DisableWarn,
            LogLevel::Info => LogFlag::DisableInfo,
            LogLevel::Debug => LogFlag::DisableDebug,
        }
    }
}

pub fn set_log_level_state(level: LogLevel, enabled: bool) {
    let flag = level.as_disable_flag();
    unsafe {
        if enabled {
            LOGGER.flags &= !flag
        } else {
            LOGGER.flags |= flag
        }
    }
}

fn log_with_writer<F: FnOnce(LogWriter)>(level: LogLevel, f: F) {
    let logger = unsafe { LOGGER };
    if (logger.flags & level.as_disable_flag()) != 0 {
        return;
    }
    f(logger.write);
    if matches!(level, LogLevel::ErrorCxx) && (logger.flags & LogFlag::ExitOnError) != 0 {
        exit(1);
    }
}

pub fn log_from_cxx(level: LogLevelCxx, msg: &Utf8CStr) {
    if let Some(level) = LogLevel::from_i32(level.repr) {
        log_with_writer(level, |write| write(level, msg));
    }
}

pub fn log_with_formatter<F: FnOnce(Formatter) -> fmt::Result>(level: LogLevel, f: F) {
    log_with_writer(level, |write| {
        let mut buf = Utf8CStrBufArr::default();
        f(&mut buf).ok();
        write(level, &buf);
    });
}

pub fn log_with_args(level: LogLevel, args: Arguments) {
    log_with_formatter(level, |w| w.write_fmt(args));
}

pub fn cmdline_logging() {
    fn cmdline_write(level: LogLevel, msg: &Utf8CStr) {
        if matches!(level, LogLevel::Info) {
            stdout().write_all(msg.as_bytes()).ok();
        } else {
            stderr().write_all(msg.as_bytes()).ok();
        }
    }

    let logger = Logger {
        write: cmdline_write,
        flags: LogFlag::ExitOnError,
    };
    unsafe {
        LOGGER = logger;
    }
}

#[macro_export]
macro_rules! error {
    ($($args:tt)+) => {
        $crate::log_with_args($crate::LogLevel::Error, format_args_nl!($($args)+))
    }
}

#[macro_export]
macro_rules! warn {
    ($($args:tt)+) => {
        $crate::log_with_args($crate::LogLevel::Warn, format_args_nl!($($args)+))
    }
}

#[macro_export]
macro_rules! info {
    ($($args:tt)+) => {
        $crate::log_with_args($crate::LogLevel::Info, format_args_nl!($($args)+))
    }
}

#[cfg(debug_assertions)]
#[macro_export]
macro_rules! debug {
    ($($args:tt)+) => {
        $crate::log_with_args($crate::LogLevel::Debug, format_args_nl!($($args)+))
    }
}

#[cfg(not(debug_assertions))]
#[macro_export]
macro_rules! debug {
    ($($args:tt)+) => {};
}

#[derive(Default)]
pub struct LoggedError {}

// Automatically handle all printable errors
impl<T: Display> From<T> for LoggedError {
    #[cfg(not(debug_assertions))]
    fn from(e: T) -> Self {
        log_with_args(LogLevel::Error, format_args_nl!("{:#}", e));
        LoggedError::default()
    }

    #[track_caller]
    #[cfg(debug_assertions)]
    fn from(e: T) -> Self {
        let caller = Location::caller();
        log_with_args(
            LogLevel::Error,
            format_args_nl!("[{}:{}] {:#}", caller.file(), caller.line(), e),
        );
        LoggedError::default()
    }
}

pub type LoggedResult<T> = Result<T, LoggedError>;

#[macro_export]
macro_rules! log_err {
    ($msg:literal $(,)?) => {{
        $crate::log_with_args($crate::LogLevel::Error, format_args_nl!($msg));
        $crate::LoggedError::default()
    }};
    ($err:expr $(,)?) => {{
        $crate::log_with_args($crate::LogLevel::Error, format_args_nl!("{}", $err));
        $crate::LoggedError::default()
    }};
    ($($args:tt)+) => {{
        $crate::log_with_args($crate::LogLevel::Error, format_args_nl!($($args)+));
        $crate::LoggedError::default()
    }};
}

pub trait ResultExt<T> {
    fn log(self) -> LoggedResult<T>;
    fn log_with_msg<F: FnOnce(Formatter) -> fmt::Result>(self, f: F) -> LoggedResult<T>;
}

pub trait ResultNoLog<T> {
    fn no_log(self) -> LoggedResult<T>;
}

// Internal C++ bridging logging routines
pub(crate) trait CxxResultExt<T> {
    fn log_cxx(self) -> LoggedResult<T>;
    fn log_cxx_with_msg<F: FnOnce(Formatter) -> fmt::Result>(self, f: F) -> LoggedResult<T>;
}

trait LogImpl<T> {
    fn log_impl(self, level: LogLevel, caller: Option<&'static Location>) -> LoggedResult<T>;
    fn log_with_msg_impl<F: FnOnce(Formatter) -> fmt::Result>(
        self,
        level: LogLevel,
        caller: Option<&'static Location>,
        f: F,
    ) -> LoggedResult<T>;
}

impl<T, E> ResultNoLog<T> for Result<T, E> {
    fn no_log(self) -> LoggedResult<T> {
        match self {
            Ok(v) => Ok(v),
            Err(_) => Err(LoggedError::default()),
        }
    }
}

impl<T> ResultNoLog<T> for Option<T> {
    fn no_log(self) -> LoggedResult<T> {
        match self {
            Some(v) => Ok(v),
            None => Err(LoggedError::default()),
        }
    }
}

impl<T, R: LogImpl<T>> CxxResultExt<T> for R {
    fn log_cxx(self) -> LoggedResult<T> {
        self.log_impl(LogLevel::ErrorCxx, None)
    }

    fn log_cxx_with_msg<F: FnOnce(Formatter) -> fmt::Result>(self, f: F) -> LoggedResult<T> {
        self.log_with_msg_impl(LogLevel::ErrorCxx, None, f)
    }
}

impl<T, R: LogImpl<T>> ResultExt<T> for R {
    #[cfg(not(debug_assertions))]
    fn log(self) -> LoggedResult<T> {
        self.log_impl(LogLevel::Error, None)
    }

    #[track_caller]
    #[cfg(debug_assertions)]
    fn log(self) -> LoggedResult<T> {
        self.log_impl(LogLevel::Error, Some(Location::caller()))
    }

    #[cfg(not(debug_assertions))]
    fn log_with_msg<F: FnOnce(Formatter) -> fmt::Result>(self, f: F) -> LoggedResult<T> {
        self.log_with_msg_impl(LogLevel::Error, None, f)
    }

    #[track_caller]
    #[cfg(debug_assertions)]
    fn log_with_msg<F: FnOnce(Formatter) -> fmt::Result>(self, f: F) -> LoggedResult<T> {
        self.log_with_msg_impl(LogLevel::Error, Some(Location::caller()), f)
    }
}

impl<T> LogImpl<T> for LoggedResult<T> {
    fn log_impl(self, _: LogLevel, _: Option<&'static Location>) -> LoggedResult<T> {
        self
    }

    fn log_with_msg_impl<F: FnOnce(Formatter) -> fmt::Result>(
        self,
        level: LogLevel,
        caller: Option<&'static Location>,
        f: F,
    ) -> LoggedResult<T> {
        match self {
            Ok(v) => Ok(v),
            Err(_) => {
                log_with_formatter(level, |w| {
                    if let Some(caller) = caller {
                        write!(w, "[{}:{}] ", caller.file(), caller.line())?;
                    }
                    f(w)?;
                    w.write_char('\n')
                });
                Err(LoggedError::default())
            }
        }
    }
}

impl<T, E: Display> LogImpl<T> for Result<T, E> {
    fn log_impl(self, level: LogLevel, caller: Option<&'static Location>) -> LoggedResult<T> {
        match self {
            Ok(v) => Ok(v),
            Err(e) => {
                if let Some(caller) = caller {
                    log_with_args(
                        level,
                        format_args_nl!("[{}:{}] {:#}", caller.file(), caller.line(), e),
                    );
                } else {
                    log_with_args(level, format_args_nl!("{:#}", e));
                }
                Err(LoggedError::default())
            }
        }
    }

    fn log_with_msg_impl<F: FnOnce(Formatter) -> fmt::Result>(
        self,
        level: LogLevel,
        caller: Option<&'static Location>,
        f: F,
    ) -> LoggedResult<T> {
        match self {
            Ok(v) => Ok(v),
            Err(e) => {
                log_with_formatter(level, |w| {
                    if let Some(caller) = caller {
                        write!(w, "[{}:{}] ", caller.file(), caller.line())?;
                    }
                    f(w)?;
                    writeln!(w, ": {:#}", e)
                });
                Err(LoggedError::default())
            }
        }
    }
}