// Copyright 2021 Brian J. Tarricone <brian@tarricone.org>
//
// This file is part of ESPMonitor.
//
// ESPMonitor is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// ESPMonitor is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with ESPMonitor.  If not, see <https://www.gnu.org/licenses/>.

use addr2line::Context;
use crossterm::{
    QueueableCommand,
    event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
    style::{Color, Print, PrintStyledContent, Stylize},
    terminal::{disable_raw_mode, enable_raw_mode},
};
use gimli::{read::Reader, EndianRcSlice, RunTimeEndian};
use lazy_static::lazy_static;
use object::read::Object;
use regex::Regex;
use serial::{self, BaudRate, SerialPort, SystemPort};
use std::{
    fs,
    io::{self, ErrorKind, Read, Write, stdout},
    process::exit,
    time::{Duration, Instant},
};

mod types;

pub use types::{AppArgs, Chip, Framework};

const DEFAULT_BAUD_RATE: BaudRate = BaudRate::Baud115200;
const UNFINISHED_LINE_TIMEOUT: Duration = Duration::from_secs(5);

lazy_static! {
    static ref LINE_SEP_RE: Regex = Regex::new("\r?\n")
        .expect("Failed to parse line separator regex");
    static ref FUNC_ADDR_RE: Regex = Regex::new(r"0x4[0-9a-fA-F]{7}")
        .expect("Failed to parse program address regex");
}

macro_rules! rprintln {
    () => (print!("\r\n"));
    ($fmt:literal) => (print!(concat!($fmt, "\r\n")));
    ($fmt:literal, $($arg:tt)+) => (print!(concat!($fmt, "\r\n"), $($arg)*));
}

struct Symbols<'a> {
    obj: object::read::File<'a, &'a [u8]>,
    context: Context<EndianRcSlice<RunTimeEndian>>,
}

struct SerialState<'a> {
    unfinished_line: String,
    last_unfinished_line_at: Instant,
    symbols: Option<Symbols<'a>>,
}

#[cfg(unix)]
pub fn run(args: AppArgs) -> Result<(), Box<dyn std::error::Error>> {
    use nix::{sys::wait::{WaitStatus, waitpid}, unistd::{ForkResult, fork}};

    enable_raw_mode()?;

    match unsafe { fork() } {
        Err(err) => {
            disable_raw_mode()?;
            Err(err.into())
        },
        Ok(ForkResult::Parent { child }) => loop {
            match waitpid(child, None) {
                Ok(WaitStatus::Exited(_, status)) => {
                    disable_raw_mode()?;
                    exit(status);
                },
                Ok(WaitStatus::Signaled(_, _, _)) => {
                    disable_raw_mode()?;
                    exit(255);
                },
                _ => (),
            }
        },
        Ok(ForkResult::Child) => run_child(args),
    }
}

#[cfg(windows)]
pub fn run(args: AppArgs) -> Result<(), Box<dyn std::error::Error>> {
    enable_raw_mode()?;
    let result = run_child(args);
    disable_raw_mode()?;
    result
}

fn run_child(args: AppArgs) -> Result<(), Box<dyn std::error::Error>> {
    rprintln!("ESPMonitor {}", env!("CARGO_PKG_VERSION"));
    rprintln!();
    rprintln!("Commands:");
    rprintln!("    CTRL+R    Reset chip");
    rprintln!("    CTRL+C    Exit");
    rprintln!();

    let speed = args.speed.map(BaudRate::from_speed).unwrap_or(DEFAULT_BAUD_RATE);
    rprintln!("Opening {} with speed {}", args.serial, speed.speed());

    let mut dev = serial::open(&args.serial)?;
    dev.set_timeout(Duration::from_millis(200))?;
    dev.reconfigure(&|settings| {
        settings.set_baud_rate(speed)
    })?;

    let bin_data = args.bin.as_ref().and_then(|bin_name| match fs::read(bin_name) {
        Ok(bin_data) => {
            rprintln!("Using {} as flash image", bin_name.to_string_lossy());
            Some(bin_data)
        },
        Err(err) => {
            rprintln!("WARNING: Unable to open flash image {}: {}", bin_name.to_string_lossy(), err);
            None
        },
    });

    let symbols = bin_data.as_ref().and_then(|bin_data| match load_bin_context(bin_data.as_slice()) {
        Ok(symbols) => Some(symbols),
        Err(err) => {
            rprintln!("WARNING: Failed to parse flash image: {}", err);
            None
        },
    });

    if args.reset {
        reset_chip(&mut dev)?;
    }


    let mut serial_state = SerialState {
        unfinished_line: String::new(),
        last_unfinished_line_at: Instant::now(),
        symbols,
    };

    let mut output = stdout();
    let mut buf = [0u8; 1024];
    loop {
        match dev.read(&mut buf) {
            Ok(bytes) if bytes > 0 => handle_serial(&mut serial_state, &buf[0..bytes], &mut output)?,
            Ok(_) => if dev.read_dsr().is_err() {
                rprintln!("Device disconnected; exiting");
                break Ok(());
            },
            Err(err) if err.kind() == ErrorKind::TimedOut => (),
            Err(err) if err.kind() == ErrorKind::WouldBlock => (),
            Err(err) if err.kind() == ErrorKind::Interrupted => (),
            Err(err) => break Err(err.into()),
        }

        while event::poll(Duration::ZERO)? {
            match event::read() {
                Ok(Event::Key(key_event)) => handle_input(&mut dev, key_event)?,
                Ok(_) => (),
                Err(err) => return Err(err.into()),
            }
        }
    }
}

fn load_bin_context(data: &[u8]) -> Result<Symbols, Box<dyn std::error::Error + 'static>> {
    let obj = object::File::parse(data)?;
    let context = Context::new(&obj)?;
    Ok(Symbols {
        obj,
        context,
    })
}

fn reset_chip(dev: &mut SystemPort) -> io::Result<()> {
    print!("Resetting device... ");
    std::io::stdout().flush()?;
    dev.set_dtr(false)?;
    dev.set_rts(true)?;
    dev.set_rts(false)?;
    rprintln!("done");
    Ok(())
}

fn handle_serial(state: &mut SerialState, buf: &[u8], output: &mut dyn Write) -> io::Result<()> {
    let data = String::from_utf8_lossy(buf);
    let mut lines = LINE_SEP_RE.split(&data).collect::<Vec<&str>>();

    let new_unfinished_line =
        if data.ends_with('\n') {
            None
        } else {
            lines.pop()
        };

    for line in lines {
        let full_line =
            if !state.unfinished_line.is_empty() {
                state.unfinished_line.push_str(line);
                state.unfinished_line.as_str()
            } else {
                line
            };

        if !full_line.is_empty() {
            output_line(state, full_line, output)?;
            state.unfinished_line.clear();
        }
    }

    if let Some(nel) = new_unfinished_line {
        state.unfinished_line.push_str(nel);
        state.last_unfinished_line_at = Instant::now();
    } else if !state.unfinished_line.is_empty() && state.last_unfinished_line_at.elapsed() > UNFINISHED_LINE_TIMEOUT {
        output_line(state, &state.unfinished_line, output)?;
        state.unfinished_line.clear();
    }

    Ok(())
}

fn output_line(state: &SerialState, line: &str, output: &mut dyn Write) -> io::Result<()> {
    if let Some(symbols) = state.symbols.as_ref() {
        let mut cur_start_offset = 0;

        for mat in FUNC_ADDR_RE.find_iter(line) {
            let (function, file, lineno) = u64::from_str_radix(&mat.as_str()[2..], 16)
                .ok()
                .map(|addr| {
                    let function = find_function_name(symbols, addr);
                    let (file, lineno) = find_location(symbols, addr);
                    (function, file, lineno)
                })
                .unwrap_or((None, None, None));

            fn or_qq(s: Option<String>) -> String {
                s.unwrap_or_else(|| "??".to_string())
            }

            output.queue(Print(line[cur_start_offset..mat.end()].to_string()))?;
            cur_start_offset = mat.end();
            let symbolicated_name = format!(" [{}:{}:{}]", or_qq(function), or_qq(file), or_qq(lineno.map(|l| l.to_string())))
                .with(Color::Yellow);
            output.queue(PrintStyledContent(symbolicated_name))?;
        }

        if cur_start_offset < line.len() {
            output.queue(Print(line[cur_start_offset..line.len()].to_string()))?;
        }
    } else {
        output.queue(Print(line.to_string()))?;
    }

    output.write_all(b"\r\n")?;
    output.flush()?;

    Ok(())
}

fn handle_input(dev: &mut SystemPort, key_event: KeyEvent) -> io::Result<()> {
    if key_event.modifiers == KeyModifiers::CONTROL {
        match key_event.code {
            KeyCode::Char('r') => reset_chip(dev),
            KeyCode::Char('c') => exit(0),
            _ => Ok(()),
        }
    } else {
        Ok(())
    }
}

fn find_function_name(symbols: &Symbols<'_>, addr: u64) -> Option<String> {
    symbols.context
        .find_frames(addr)
        .ok()
        .and_then(|mut frames| frames.next().ok().flatten())
        .and_then(|frame| frame.function.and_then(|function| function.name.to_string_lossy().ok().map(|name| name.into_owned())))
        .or_else(|| symbols.obj.symbol_map().get(addr).map(|sym| sym.name().to_string()))
}

fn find_location(symbols: &Symbols<'_>, addr: u64) -> (Option<String>, Option<u32>) {
    symbols.context
        .find_location(addr)
        .ok()
        .map(|location| (
            location.as_ref().and_then(|location| location.file).map(|file| file.to_string()),
            location.as_ref().and_then(|location| location.line)
        ))
        .unwrap_or((None, None))
}