pub mod disassembly;

use pest::iterators::{Pair, Pairs};
use pest::Parser;

use crate::test_binary::{RelayAction, RelayInstr, Segment, SensorAction, SensorInstr, Test};
use crate::{result::CompilerResult, test_binary::DeviceAddress};

// The following derive macro links the "assembly.pest" file to the Parser
#[derive(Parser)]
#[grammar = "./test_assembly/assembly.pest"]
struct TAFParser;

/// Parses a test assembly program into a Test if possible
pub fn parse(text: String) -> CompilerResult<Test> {
    let mut res = CompilerResult::new("Assembly Parsing");

    // Parse the input text using the Pest parser
    match TAFParser::parse(Rule::TEST, &text) {
        // Parse the resulting tree
        Ok(parse_tree) => {
            res.set(test_from_tree(parse_tree.peek().unwrap()));
        }
        // Record the parse error
        Err(pest_error) => {
            res.error(Box::new(pest_error));
        }
    }

    res
}

/// Create a Test struct from a parse tree node of type Rule::TEST
fn test_from_tree(node: Pair<Rule>) -> CompilerResult<Test> {
    let mut res = CompilerResult::unnamed();
    let mut test = Test::new();
    // The index of the current abort
    let mut current_abort: usize = 0;

    for node in node.into_inner() {
        match node.as_rule() {
            // Generate an abort from the ABORT rule
            Rule::ABORT => {
                // Try to interpret the inside of the ABORT as segments
                let mut iterator = node.into_inner();
                let abort_header = iterator.next().unwrap();
                let abort_idx_rule = abort_header.into_inner().next().unwrap();
                let abort_idx_num = abort_idx_rule.as_str().parse::<u16>().unwrap();
                if abort_idx_num != current_abort as u16 + 1 {
                    res.error(format!(
                        "The provided index of the {}th abort was {} instead!",
                        current_abort + 1,
                        abort_idx_num
                    ));
                }
                if let Some(mut segments) = res.require(segments_from_tree(iterator)) {
                    // If this would be abort 15 or less, then add the segments to the abort
                    if current_abort < 15 {
                        test.aborts[current_abort].segments.append(&mut segments);
                    }
                }
                // Increment the id of the current_abort
                current_abort += 1;
            }
            // Parse the segments for the test body
            Rule::TEST_BODY => {
                // Try to interpret the inside of the TEST_BODY as segments
                let mut iterator = node.into_inner();
                // ignore the header
                let _header = iterator.next();
                if let Some(mut segments) = res.require(segments_from_tree(iterator)) {
                    // Add the segments to the body
                    test.body.segments.append(&mut segments);
                }
            }
            // Do nothing for the End-Of-Input rule
            Rule::EOI => {}
            // No other
            _ => unreachable!(),
        }
    }

    // If too many Aborts were found, then fail the parse
    if current_abort >= 15 {
        res.error(format!("Found {} Aborts, Limit is 15", current_abort))
    }

    res.with_value(test)
}

/// Create a sequence of Segments from a parse tree node of type Rule::TEST_BODY or Rule::ABORT
fn segments_from_tree(iterator: Pairs<Rule>) -> CompilerResult<Vec<Segment>> {
    let mut res = CompilerResult::unnamed();
    // let mut iterator = node.into_inner();
    // Ignore the header
    // let _header = iterator.next();

    res.set_value(Vec::new());
    res.collect(iterator.map(segment_from_tree));
    res
}

/// Create a segment from a parse tree of type Rule::RELAY_STATEMENT or type Rule::SENSOR_STATEMENT
fn segment_from_tree(node: Pair<Rule>) -> CompilerResult<Segment> {
    let mut res = CompilerResult::unnamed();
    // let span = node.as_span();
    let mut iterator = node.into_inner();

    #[derive(PartialEq, Eq)]
    enum State {
        Sensors,
        Relays,
    }
    // The current state of segment parsing
    // Statements within a Segment go in the order Sensors then Relays
    let mut state = State::Sensors;

    let rel_time = if let Some(segment_header) = iterator.next() {
        if let Some(relative_time) = segment_header.into_inner().next() {
            res.require(relative_time.as_str().parse::<u32>())
        } else {
            res.error("Relative Time Not Found");
            None
        }
    } else {
        res.error("Segment Header Not Found");
        None
    };

    let rel_time = if let Some(rel_time) = rel_time {
        rel_time
    } else {
        1
    };
    let mut segment = check!(res, Segment::try_new(rel_time));

    for child in iterator {
        match child.as_rule() {
            Rule::SENSOR_STATEMENT => {
                // If we have already parsed a Relay, fail
                if state == State::Relays {
                    res.error("Sensor statement found mixed in with or after Relays");
                }
                // If we can parse a sensor, add it to the sensors list
                if let Some(sensor_stmt) = res.require(sensor_statement_from_tree(child)) {
                    segment.sensor_instructions.push(sensor_stmt)
                }
            }
            Rule::RELAY_STATEMENT => {
                // Mark the state as parsing relays
                state = State::Relays;

                // If we can parse a relay, add it to the relays list
                if let Some(relay_stmt) = res.require(relay_statement_from_tree(child)) {
                    segment.relay_instructions.push(relay_stmt)
                }
            }
            _ => unreachable!(),
        }
    }

    res.with_value(segment)
}

/// Create a device address from a parse tree node of type Rule::ADDRESS
fn device_address_from_tree(node: Pair<Rule>) -> CompilerResult<DeviceAddress> {
    let mut res = CompilerResult::new("Parsing device address");

    let mut iterator = node.into_inner();

    let prefix = iterator.next().unwrap().as_str();
    let address = check!(
        res,
        iterator.next().unwrap().as_str().parse::<u16>(),
        "Invalid device address"
    );

    let virtuality = match prefix {
        "P" => 0,
        "V" => 1,
        _ => {
            res.error("Invalid prefix for device address");
            return res;
        }
    };

    res.with_value(DeviceAddress {
        value: address,
        virtuality,
    })
}

/// Create a Relay Statement from a parse tree node of type Rule::RELAY_STATEMENT
fn relay_statement_from_tree(node: Pair<Rule>) -> CompilerResult<RelayInstr> {
    let mut res = CompilerResult::new("Parsing Relay Statement");

    let mut iterator = node.into_inner();
    let op = iterator.next().unwrap();

    let address = check!(res, device_address_from_tree(iterator.next().unwrap()));

    let action = check!(res, RelayAction::parse(&op.as_str().to_lowercase()));

    let relay_instr = check!(res, RelayInstr::try_new(action, address));
    res.with_value(relay_instr)
}

/// Create a Sensor Statement from a parse tree node of type Rule::SENSOR_STATEMENT
fn sensor_statement_from_tree(node: Pair<Rule>) -> CompilerResult<SensorInstr> {
    let inner = node.into_inner().next().unwrap();

    match inner.as_rule() {
        Rule::SENSOR_CHECK_STATEMENT => sensor_check_from_tree(inner),
        Rule::SENSOR_STOP_STATEMENT => sensor_stop_from_tree(inner),
        _ => unreachable!(),
    }
}

/// Create a Sensor Check from a parse tree node of type SENSOR_CHECK_STATEMENT
fn sensor_check_from_tree(node: Pair<Rule>) -> CompilerResult<SensorInstr> {
    let mut res = CompilerResult::new("Parsing Sensor Check Instruction");
    // let span = node.as_span();
    let mut iterator = node.into_inner();

    // Retrieve the various nodes
    let sensor_check = iterator.next().unwrap();
    let address = check!(res, device_address_from_tree(iterator.next().unwrap()));
    let left_bound = iterator.next().unwrap();
    let right_bound = iterator.next().unwrap();
    let abort_idx = iterator.next().unwrap();

    // Parse the various sub-trees
    let action = check!(
        res,
        SensorAction::parse(&sensor_check.as_str().to_lowercase())
    );
    let left_bound = check!(res, parse_hex(left_bound));
    let right_bound = check!(res, parse_hex(right_bound));
    let abort_idx = check!(res, abort_idx.as_str().parse::<u8>());
    //Create a new sensor instruction from the parsed values
    let sensor_instr = check!(
        res,
        SensorInstr::try_new(action, abort_idx, address, left_bound, right_bound,)
    );
    res.with_value(sensor_instr)
}

/// Create a Sensor Stop from a parse tree node of type SENSOR_STOP_STATEMENT
fn sensor_stop_from_tree(node: Pair<Rule>) -> CompilerResult<SensorInstr> {
    let mut res = CompilerResult::new("Parsing Sensor Stop Instruction");
    // let span = node.as_span();
    let mut iterator = node.into_inner();

    let address = check!(res, device_address_from_tree(iterator.next().unwrap()));
    let sensor_instr = check!(res, SensorInstr::try_new_stop(address));
    res.with_value(sensor_instr)
}

fn parse_hex(node: Pair<Rule>) -> CompilerResult<u16> {
    let mut res = CompilerResult::new("Parsing Hexadecimal Number Literal");
    // let span = node.as_span();
    res.set_res(u16::from_str_radix(
        node.as_str().trim_start_matches("0x"),
        16,
    ));
    res
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::result::Status;

    #[test]
    fn test_valid_hex_parse() {
        let valid_hex = "0xAA";

        let parsed_hex = TAFParser::parse(Rule::HEX_NUM, valid_hex);
        let hex_val = parse_hex(parsed_hex.unwrap().peek().unwrap());
        assert_eq!(Some(170), hex_val.to_option());
    }

    #[test]
    fn test_invalid_hex_parse() {
        let valid_hex = "0xZAZA";
        assert_eq!(true, TAFParser::parse(Rule::HEX_NUM, valid_hex).is_err());
    }

    #[test]
    fn test_abort_header_parse() {
        for i in 1..16 {
            let abort_header_1 = format!("abort #{}:\n", i);
            assert_eq!(
                true,
                TAFParser::parse(Rule::ABORT_HEADER, abort_header_1.as_str()).is_ok()
            )
        }
    }

    #[test]
    fn test_abort_header_idx_check() {
        let test_valid = r#"abort #1:
        segment +1ms

        abort #2:
        segment +1ms
        
        abort #3:
        segment +1ms
        
        test:
        segment +1ms
        "#;
        let test_invalid = r#"abort #1:
        segment +1ms

        abort #3:
        segment +1ms
        
        abort #2:
        segment +1ms
        
        test:
        segment +1ms
        "#;
        assert_eq!(Status::Passed, parse(test_valid.to_string()).get_status());
        assert_eq!(Status::Failed, parse(test_invalid.to_string()).get_status());
    }

    #[test]
    fn test_sensor_check_instant() {
        let sensor_check = "is_now_in P1, 0x00, 0xFF, #1\n";
        let expected = SensorInstr::try_new(
            SensorAction::is_now_in,
            1,
            DeviceAddress {
                value: 1,
                virtuality: 0,
            },
            0,
            255,
        )
        .to_option()
        .unwrap();
        let pairs = TAFParser::parse(Rule::SENSOR_CHECK_STATEMENT, sensor_check);
        let check_instr = sensor_check_from_tree(pairs.unwrap().peek().unwrap());
        assert_eq!(Some(expected), check_instr.to_option());
    }

    #[test]
    fn test_sensor_start_check() {
        let sensor_check_statement = "start_check_in P1, 0x00, 0xFF, #1\n";
        let sensor_check = SensorInstr::try_new(
            SensorAction::start_check_in,
            1,
            DeviceAddress {
                value: 1,
                virtuality: 0,
            },
            0,
            255,
        )
        .to_option()
        .unwrap();
        let sensor_pairs_start =
            TAFParser::parse(Rule::SENSOR_CHECK_STATEMENT, sensor_check_statement);
        let check_start_statement =
            sensor_check_from_tree(sensor_pairs_start.unwrap().peek().unwrap());
        assert_eq!(Some(sensor_check), check_start_statement.to_option());
    }

    #[test]
    fn test_sensor_stop_check() {
        let sensor_stop_statement = "stop_check P3\n";
        let sensor_stop = SensorInstr::try_new_stop(DeviceAddress {
            value: 3,
            virtuality: 0,
        })
        .to_option()
        .unwrap();
        let sensor_pairs_stop =
            TAFParser::parse(Rule::SENSOR_STOP_STATEMENT, sensor_stop_statement);
        let check_stop_statement =
            sensor_stop_from_tree(sensor_pairs_stop.unwrap().peek().unwrap());
        println!("{:?}", check_stop_statement);
        assert_eq!(Some(sensor_stop), check_stop_statement.to_option());
    }

    #[test]
    fn test_relay_set() {
        let relay_statement_set = "set P1\n";
        let relay_set = RelayInstr::try_new(
            RelayAction::Set,
            DeviceAddress {
                value: 1,
                virtuality: 0,
            },
        )
        .to_option()
        .unwrap();
        let relay_pairs_set = TAFParser::parse(Rule::RELAY_STATEMENT, relay_statement_set);
        let check_set_statement =
            relay_statement_from_tree(relay_pairs_set.unwrap().peek().unwrap());
        assert_eq!(Some(relay_set), check_set_statement.to_option());
    }

    #[test]
    fn test_relay_unset() {
        let relay_statement_unset = "unset P1\n";
        let relay_unset = RelayInstr::try_new(
            RelayAction::Unset,
            DeviceAddress {
                value: 1,
                virtuality: 0,
            },
        )
        .to_option()
        .unwrap();
        let relay_pairs_unset = TAFParser::parse(Rule::RELAY_STATEMENT, relay_statement_unset);
        let check_unset_statement =
            relay_statement_from_tree(relay_pairs_unset.unwrap().peek().unwrap());

        assert_eq!(Some(relay_unset), check_unset_statement.to_option())
    }

    #[test]
    fn white_space_in_middle() {
        let relay_statement = "set  P1           \n";
        let relay_expected = RelayInstr::try_new(
            RelayAction::Set,
            DeviceAddress {
                value: 1,
                virtuality: 0,
            },
        )
        .to_option()
        .unwrap();
        let relay_pairs = TAFParser::parse(Rule::RELAY_STATEMENT, relay_statement);
        let relay_instruction = relay_statement_from_tree(relay_pairs.unwrap().peek().unwrap());
        assert_eq!(Some(relay_expected), relay_instruction.to_option());

        let sensor_statement = "is_now_in               P1,            0x00,                  0xFF,           #1            \n";
        let sensor_expected = SensorInstr::try_new(
            SensorAction::is_now_in,
            1,
            DeviceAddress {
                value: 1,
                virtuality: 0,
            },
            0,
            255,
        )
        .to_option()
        .unwrap();
        let sensor_pairs = TAFParser::parse(Rule::SENSOR_CHECK_STATEMENT, sensor_statement);
        let sensor_instruction = sensor_check_from_tree(sensor_pairs.unwrap().peek().unwrap());
        assert_eq!(Some(sensor_expected), sensor_instruction.to_option());
    }

    #[test]
    fn extra_new_line_at_end() {
        let relay_statement = "set P1\n\n\n\n\n\n\n\n\n";
        let relay_expected = RelayInstr::try_new(
            RelayAction::Set,
            DeviceAddress {
                value: 1,
                virtuality: 0,
            },
        )
        .to_option()
        .unwrap();
        let relay_pairs = TAFParser::parse(Rule::RELAY_STATEMENT, relay_statement);
        let relay_instruction = relay_statement_from_tree(relay_pairs.unwrap().peek().unwrap());
        assert_eq!(Some(relay_expected), relay_instruction.to_option());

        let sensor_statement = "is_now_in P1, 0x00, 0xFF, #1\n\n\n\n\n\n\n\n\n";
        let sensor_expected = SensorInstr::try_new(
            SensorAction::is_now_in,
            1,
            DeviceAddress {
                value: 1,
                virtuality: 0,
            },
            0,
            255,
        )
        .to_option()
        .unwrap();
        let sensor_pairs = TAFParser::parse(Rule::SENSOR_CHECK_STATEMENT, sensor_statement);
        let sensor_instruction = sensor_check_from_tree(sensor_pairs.unwrap().peek().unwrap());
        assert_eq!(Some(sensor_expected), sensor_instruction.to_option());
    }
    #[test]
    fn new_line_in_middle() {
        let relay_statement = "set \n\n\n\n P1\n";
        assert!(TAFParser::parse(Rule::RELAY_STATEMENT, relay_statement).is_err());

        let sensor_statement = "is_now_in P1 \n\n, 0x00\n, 0xFF\n, #1\n";
        assert!(TAFParser::parse(Rule::SENSOR_CHECK_STATEMENT, sensor_statement).is_err());
    }

    #[test]
    fn new_line_in_start() {
        let relay_statement = "\n\n\n\nset P1\n";
        assert!(TAFParser::parse(Rule::RELAY_STATEMENT, relay_statement).is_err());

        let sensor_statement = "\n\n\n\nis_now_in 1, 0x00, 0xFF, P1\n";
        assert!(TAFParser::parse(Rule::SENSOR_CHECK_STATEMENT, sensor_statement).is_err());
    }
}