use core::cell::RefCell;
use core::fmt::Write;

use alloc::{string::String, vec::Vec};
use backtrace::{BytesOrWideString, Symbol, SymbolName};
use critical_section::{CriticalSection, Mutex};
use esp_hal::uart::UartTx;
use esp_hal::Blocking;
use log::{LevelFilter, Log};

static ALT_LOGGER_UART: Mutex<RefCell<Option<UartTx<'static, Blocking>>>> = Mutex::new(RefCell::new(None));

pub fn with_uart_tx<R>(f: impl FnOnce(CriticalSection<'_>, &'_ mut UartTx<'static, Blocking>) -> R) -> R {
    critical_section::with(|cs| {
        let mut uart = ALT_LOGGER_UART.borrow(cs).borrow_mut();
        let uart = uart.as_mut().unwrap();

        (f)(cs, uart)
    })
}

struct AlternativeLogger;

impl Log for AlternativeLogger {
    #[allow(unused)]
    fn enabled(&self, _: &log::Metadata) -> bool {
        // Filtered by `log` already
        true
    }

    #[allow(unused)]
    fn log(&self, record: &log::Record) {
        with_uart_tx(|cs, uart| {
            print_log_record(uart, record);
        })
    }

    fn flush(&self) {}
}

const RESET: &str = "\u{001B}[0m";
const RED: &str = "\u{001B}[31m";
const GREEN: &str = "\u{001B}[32m";
const YELLOW: &str = "\u{001B}[33m";
const BLUE: &str = "\u{001B}[34m";
const CYAN: &str = "\u{001B}[35m";

#[cfg(feature = "alt-log")]
macro_rules! println {
    () => {{
        do_print(Default::default());
    }};

    ($($arg:tt)*) => {{
        do_print(::core::format_args!($($arg)*));
    }};
}


fn do_print(args: core::fmt::Arguments<'_>) {
    with_uart_tx(|_, uart| {
        uart.write_fmt(args).unwrap();
        uart.write_str("\n").unwrap();
        uart.flush().unwrap();
    })
}

fn print_log_record(uart: &mut UartTx<'_, Blocking>, record: &log::Record) {
    let color = match record.level() {
        log::Level::Error => RED,
        log::Level::Warn => YELLOW,
        log::Level::Info => GREEN,
        log::Level::Debug => BLUE,
        log::Level::Trace => CYAN,
    };
    let reset = RESET;
    let args = format_args!("{}{:>5} - {}{}\n", color, record.level(), record.args(), reset);
    uart.write_fmt(args).unwrap();
    uart.flush().unwrap();
}

#[cfg(not(feature = "usb-log"))]
#[panic_handler]
fn panic_handler(info: &core::panic::PanicInfo) -> ! {
    use esp_backtrace::Backtrace;

    println!("{RED}");
    println!("=============== CUSTOM PANIC HANDLER ==============");
    println!("{info}{RESET}");
    println!("");
    println!("Backtrace:");
    println!("");

    let backtrace = Backtrace::capture();

    for frame in backtrace.frames() {
        println!("0x{:x}", frame.program_counter());
        print_resolved_symbol(frame.program_counter());
    }

    loop {}
}

#[cfg(not(feature = "usb-log"))]
fn print_resolved_symbol(address: usize) {
    unsafe {
        backtrace::resolve_unsynchronized(address as _, |sym| {
            println!("? {:x?}\n{}\n\n", sym.addr().map(|addr| addr as usize), sym.name().unwrap_or(SymbolName::new(b"(unknown)")));
            // sym.filename_raw().unwrap_or(BytesOrWideString::Bytes(&[]));
            // let segments: Vec<String> = vec![
            //     sym.addr().map(|a| format!("{:x?}", a as isize)),
            //     sym.name().map(|n| n.to_string()),
            //     sym.filename().map(|p| {
            //         format!(
            //             "{}{}",
            //             p.to_string_lossy(),
            //             sym.lineno()
            //                 .map(|n| format!(" (line {})", n))
            //                 .unwrap_or_else(|| "".to_string())
            //         )
            //     }),
            // ]
            // .into_iter()
            // .flatten()
            // .collect();
            // segments.join("\n")
        });
    }
}

pub fn setup_alternative_logging(alt_uart: UartTx<'static, Blocking>, level_filter: LevelFilter)
{
    critical_section::with(|cs| {
        *ALT_LOGGER_UART.borrow(cs).borrow_mut() = Some(alt_uart);
    });

    unsafe {
        log::set_logger_racy(&AlternativeLogger).unwrap();
        log::set_max_level_racy(level_filter);
    }
}