use std::cmp::max;
use std::collections::{HashMap, HashSet};
use std::fs::read_to_string;
use std::path::Path;
use std::path::PathBuf;
use std::str::FromStr;

use crate::environment::config::{
    Config, RelayRecord, SensorRecord, Virtual, VirtualRelayRecord, VirtualSensorRecord,
};
use crate::environment::drivers::Drivers;
use crate::environment::{collect_calibration_tables, Environment};
use crate::result::diagnostic::Location;
use crate::test_descriptor::ast::*;

/// Helper function to construct an empty diagnostic location
fn empty_location() -> Location {
    Location::from_raw(0, 0, 1)
}

/// Builder for a TDF test struct and the associated environment
#[derive(Debug)]
pub struct TestBuilder {
    pub test: Test,
    pub env: Environment,
}

/// Builder for a vector of instructions for either a test body or abort
pub struct InstructionStreamBuilder<'a> {
    pub test_builder: &'a mut TestBuilder,
    pub stream_id: StreamId<'a>,
}

/// Identify instruction streams as either within an abort or test body
#[derive(Clone)]
pub enum StreamId<'a> {
    Abort(usize),
    Body,
    Section {
        name: String,
        stream_id: &'a StreamId<'a>,
    },
}

/// Builder for a single instruction at an instant relative time
pub struct MomentBuilder<'a> {
    pub test_builder: &'a mut TestBuilder,
    pub stream_id: &'a StreamId<'a>,
    pub time: u128,
}

/// Builder for a single instruction over an interval of relative time
pub struct IntervalBuilder<'a> {
    pub test_builder: &'a mut TestBuilder,
    pub stream_id: &'a StreamId<'a>,
    pub start: u128,
    pub end: u128,
}

impl TestBuilder {
    /// Create an empty TestBuilder with default test and environment
    pub fn new() -> Self {
        TestBuilder {
            test: TestBuilder::default_test(),
            env: TestBuilder::default_env(),
        }
    }

    /// Build a completely empty Test struct
    fn default_test() -> Test {
        let test_body = TestBody {
            statements: Vec::new(),
        };
        let aborts = Vec::new();
        Test { test_body, aborts }
    }

    /// Build a completely empty Environment struct apart from a couple necessary config values
    fn default_env() -> Environment {
        // Construct default config
        let virtuals = Virtual {
            sensors: HashMap::new(),
            relays: HashMap::new(),
        };
        let config = Config {
            relay_micros: 10u32,
            sensor_micros: 10u32,
            max_segment_length_micros: 1000u32,
            safe_state: Vec::new(),
            drivers_path: None,
            relays: HashMap::new(),
            sensors: HashMap::new(),
            virtuals,
            global_bounds: None,
        };

        // Load basic drivers file
        let drivers_filename = "test_builder_drivers.toml";
        let drivers_path = PathBuf::from(format!("tests/input_files/{}", drivers_filename));
        let drivers_contents =
            read_to_string(drivers_path).expect("Failed to read test builder drivers file");
        let drivers_file = Drivers::parse(0, drivers_filename.to_string(), &drivers_contents)
            .to_option()
            .expect("Failed to parse test builder drivers file");

        // Load default calibration file
        let table_cache = collect_calibration_tables(
            Path::new(""),
            &HashSet::from([PathBuf::from("tests/input_files/calibration_numeric.ncf")]),
        )
        .to_option()
        .unwrap();

        // Construct and return the default environment
        Environment {
            config,
            drivers_file: Some(drivers_file),
            table_cache,
        }
    }

    /// Add an abort to the test struct and return an InstructionStreamBuilder that will build
    /// instructions for this abort
    pub fn with_abort(&mut self, abort_name: &str) -> InstructionStreamBuilder {
        self.test.aborts.push(Abort {
            name: abort_name.to_string(),
            statements: Vec::new(),
        });
        let abort_idx = self.test.aborts.len() - 1;
        InstructionStreamBuilder::new(self, StreamId::Abort(abort_idx))
    }

    /// Get an InstructionStreamBuilder that will build instructions for the test body
    pub fn with_body(&mut self) -> InstructionStreamBuilder {
        InstructionStreamBuilder::new(self, StreamId::Body)
    }

    /// Add a sensor record into the environment if one with the given name does not already exist
    /// Automatically uses a 1ms polling interval and the numeric calibration file in the
    /// tests/input_files directory
    pub fn with_sensor(&mut self, sensor_name: &str) -> &mut Self {
        // Ensure we aren't creating a duplicate sensor
        if !self.env.config.sensors.contains_key(sensor_name) {
            let sensor_record = SensorRecord {
                // Make the device address use the smallest unused address
                device_address: ((self.env.config.sensors.len() + self.env.config.relays.len())
                    as u16)
                    .into(),
                polling_interval_ms: 1,
                // Always use the same calibration file. If this becomes configurable, we will need
                // to include a call to `collect_calibration_tables()`
                calibration: PathBuf::from_str("tests/input_files/calibration_numeric.ncf")
                    .unwrap(),
            };
            self.env
                .config
                .sensors
                .insert(sensor_name.to_string(), sensor_record);
        }
        self
    }

    /// Add a relay record into the environment if one with the given name does not already exist.
    pub fn with_relay(&mut self, relay_name: &str) -> &mut Self {
        if !self.env.config.relays.contains_key(relay_name) {
            let relay_record = RelayRecord {
                // Make the device address use the smallest unused address
                device_address: ((self.env.config.sensors.len() + self.env.config.relays.len())
                    as u16)
                    .into(),
            };
            self.env
                .config
                .relays
                .insert(relay_name.to_string(), relay_record);
        }
        self
    }

    /// Add a virtual sensor record into the environment if one with the given name does not already exist.
    /// Automatically uses the driver "test_sensor_driver" with no arguments, a 1ms polling interval,
    /// and the numeric calibration file in the tests/input_files directory
    pub fn with_v_sensor(&mut self, sensor_name: &str) -> &mut Self {
        if !self.env.config.virtuals.sensors.contains_key(sensor_name) {
            let v_sensor_record = VirtualSensorRecord {
                // Make the device address use the smallest unused virtual address
                device_address: ((self.env.config.virtuals.sensors.len()
                    + self.env.config.virtuals.relays.len())
                    as u16)
                    .into(),
                driver: "test/sensor_driver".to_string(),
                args: None,
                polling_interval_ms: 1,
                calibration: PathBuf::from_str("tests/input_files/calibration_numeric.ncf").ok(),
            };
            self.env
                .config
                .virtuals
                .sensors
                .insert(sensor_name.to_string(), v_sensor_record);
        }
        self
    }

    /// Add a virtual relay record into the environment if one with the given name does not already exist.
    /// Automatically uses the driver "test_relay_driver" with no arguments.
    pub fn with_v_relay(&mut self, relay_name: &str) -> &mut Self {
        if !self.env.config.virtuals.relays.contains_key(relay_name) {
            let v_relay_record = VirtualRelayRecord {
                // Make the device address use the smallest unused virtual address
                device_address: ((self.env.config.virtuals.sensors.len()
                    + self.env.config.virtuals.relays.len())
                    as u16)
                    .into(),
                driver: "test/relay_driver".to_string(),
                args: None,
            };
            self.env
                .config
                .virtuals
                .relays
                .insert(relay_name.to_string(), v_relay_record);
        }
        self
    }

    /// Add a relay name into the configured safe state
    /// No error checks are performed
    pub fn with_safe_state_relay(&mut self, relay_name: &str) -> &mut Self {
        self.env.config.safe_state.push(relay_name.into());

        self
    }

    /// Get a reference to the statements within the section statement defined with the
    /// given name. Searches for the section name within the given StreamId
    /// This is used to add statements using the StreamId::Section Variant
    fn get_section_with_name(&self, section_name: &str, stream_id: &StreamId) -> &Vec<Statement> {
        // First, get a vector of every statement contained within the provided StreamId
        let statements = self.get_statements(stream_id);

        // Then, search through the statements for the first section with the provided name
        for statement in statements {
            if let Statement::Section(section) = statement {
                if section.name == section_name {
                    return &section.statements;
                }
            }
        }
        panic!("No section found with name {}", section_name)
    }

    /// Get a mutable reference to the statements within the section statement defined with the
    /// given name. Searches for the section name within the given StreamId
    /// This is used to add statements using the StreamId::Section Variant
    fn get_section_with_name_mut(
        &mut self,
        section_name: &str,
        stream_id: &StreamId,
    ) -> &mut Vec<Statement> {
        // First, get a mutable vector of every statement contained within the provided StreamId
        let statements = self.get_statements_mut(stream_id);

        // Then, search through the statements for the first section with the provided name
        for statement in statements.iter_mut() {
            if let Statement::Section(section) = statement {
                if section.name == section_name {
                    return &mut section.statements;
                }
            }
        }
        panic!("No section found with name {}", section_name)
    }

    /// Get the vector of statements in the given stream id
    pub fn get_statements(&self, stream_id: &StreamId) -> &Vec<Statement> {
        match stream_id {
            StreamId::Abort(abort_idx) => &self.test.aborts[*abort_idx].statements,
            StreamId::Body => &self.test.test_body.statements,
            StreamId::Section { name, stream_id } => self.get_section_with_name(name, stream_id),
        }
    }

    /// Get a mutable reference to the vector of statements in the given stream id
    pub fn get_statements_mut(&mut self, stream_id: &StreamId) -> &mut Vec<Statement> {
        match stream_id {
            StreamId::Abort(abort_idx) => &mut self.test.aborts[*abort_idx].statements,
            StreamId::Body => &mut self.test.test_body.statements,
            StreamId::Section { name, stream_id } => {
                self.get_section_with_name_mut(name, stream_id)
            }
        }
    }
}

impl<'a> InstructionStreamBuilder<'a> {
    /// Create a new stream builder which will add instructions to the given stream id
    pub fn new(test_builder: &'a mut TestBuilder, stream_id: StreamId<'a>) -> Self {
        InstructionStreamBuilder {
            test_builder,
            stream_id,
        }
    }

    /// Create a moment builder for this instruction stream at an instant in time
    pub fn at(&mut self, time: u128) -> MomentBuilder<'_> {
        MomentBuilder::new(self.test_builder, &self.stream_id, time)
    }

    /// Create a moment builder for this instruction stream at an offset after the current end of the test
    pub fn after_end(&mut self, offset: u128) -> MomentBuilder<'_> {
        MomentBuilder::new(
            self.test_builder,
            &self.stream_id,
            self.get_largest_time() + offset,
        )
    }

    /// Create an interval builder for this instruction stream from a given start and end time
    pub fn from(&mut self, start: u128, end: u128) -> IntervalBuilder<'_> {
        IntervalBuilder::new(self.test_builder, &self.stream_id, start, end)
    }

    /// Create a section statement within this instruction stream and return a new instruction
    /// stream builder to add instructions to the new section statement
    pub fn section_from(
        &mut self,
        start: u128,
        duration: u128,
        section_name: &str,
    ) -> InstructionStreamBuilder<'_> {
        // Create the new section statement
        let section_statement = SectionStatement {
            name: section_name.to_string(),
            time: SectionTime { start, duration },
            statements: Vec::new(),
            metadata: empty_location(),
        };

        // Add the section statement to the test
        self.test_builder
            .get_statements_mut(&self.stream_id)
            .push(Statement::Section(section_statement));

        // Create a new StreamId for this section statement to allow us to add statements within it
        let stream_id = StreamId::Section {
            name: section_name.to_string(),
            stream_id: &self.stream_id,
        };

        // Return an InstructionStreamBuilder using the new StreamId
        InstructionStreamBuilder {
            test_builder: self.test_builder,
            stream_id,
        }
    }

    /// Helper function to find the largest time in the instr stream. Used to place new statements at the end
    fn get_largest_time(&self) -> u128 {
        let mut largest_time = 0_u128;
        for statement in self.test_builder.get_statements(&self.stream_id).iter() {
            largest_time = max(
                largest_time,
                match statement {
                    Statement::Section(section_statement) => {
                        section_statement.time.start + section_statement.time.duration
                    }
                    Statement::Relay(relay_statement) => relay_statement.time,
                    Statement::Sensor(sensor_statement) => sensor_statement.time.get_end(),
                },
            );
        }
        largest_time
    }
}

impl<'a> MomentBuilder<'a> {
    pub fn new(test_builder: &'a mut TestBuilder, stream_id: &'a StreamId, time: u128) -> Self {
        MomentBuilder {
            test_builder,
            stream_id,
            time,
        }
    }

    /// Set the provided relay at this moment
    pub fn set(self, relay_name: &str) {
        self.test_builder
            .get_statements_mut(self.stream_id)
            .push(Statement::Relay(RelayStatement {
                time: self.time,
                id: relay_name.to_string(),
                op: RelayOp::Set,
                metadata: empty_location(),
            }));
    }

    /// Unset the provided relay at this moment
    pub fn unset(self, relay_name: &str) {
        self.test_builder
            .get_statements_mut(self.stream_id)
            .push(Statement::Relay(RelayStatement {
                time: self.time,
                id: relay_name.to_string(),
                op: RelayOp::Unset,
                metadata: empty_location(),
            }));
    }

    /// Require a sensor at this moment
    pub fn require(self, sensor_name: &str, left_bound: f64, right_bound: f64, abort_name: &str) {
        let sensor_bound = SensorBound::Numeric(SensorBoundNumeric {
            left: left_bound,
            right: right_bound,
            unit: "C".to_string(),
        });

        let constraint = SensorConstraint {
            id: sensor_name.to_string(),
            sensor_bound,
            abort: abort_name.to_string(),
            metadata: empty_location(),
        };

        self.test_builder
            .get_statements_mut(self.stream_id)
            .push(Statement::Sensor(SensorStatement {
                time: SensorTime::Instant { time: self.time },
                constraints: vec![constraint],
                abort: abort_name.to_string(),
                metadata: empty_location(),
            }));
    }
}

impl<'a> IntervalBuilder<'a> {
    pub fn new(
        test_builder: &'a mut TestBuilder,
        stream_id: &'a StreamId,
        start: u128,
        end: u128,
    ) -> Self {
        IntervalBuilder {
            test_builder,
            stream_id,
            start,
            end,
        }
    }

    /// Require a sensor over this interval
    pub fn require(self, sensor_name: &str, left_bound: f64, right_bound: f64, abort_name: &str) {
        let sensor_bound = SensorBound::Numeric(SensorBoundNumeric {
            left: left_bound,
            right: right_bound,
            unit: "C".to_string(),
        });

        let constraint = SensorConstraint {
            id: sensor_name.to_string(),
            sensor_bound,
            abort: abort_name.to_string(),
            metadata: empty_location(),
        };

        let time = SensorTime::Interval {
            start: self.start,
            end: self.end,
        };

        let statements = self.test_builder.get_statements_mut(self.stream_id);
        statements.push(Statement::Sensor(SensorStatement {
            time,
            constraints: vec![constraint],
            abort: abort_name.to_string(),
            metadata: empty_location(),
        }));
    }
}

#[cfg(test)]
mod tests {
    use std::io::Read;

    use super::*;
    use crate::test_assembly::disassembly;
    use crate::test_descriptor::concrete_test::ConcreteTest;
    use crate::timeline;
    use std::fs;
    use tempfile::NamedTempFile;

    #[test]
    fn testbuilder() {
        let mut testbuilder = TestBuilder::new();
        testbuilder.with_sensor("sensor1").with_sensor("sensor2");
        testbuilder.with_relay("relay1").with_relay("relay2");
        testbuilder.with_v_sensor("v_sensor1");
        testbuilder.with_v_relay("v_relay1");
        // let mut abort1: InstructionStreamBuilder = testbuilder.with_abort("abort1");
        // abort1.with_relay_set("relay3", 1);
        // abort1.with_relay_unset_end("relay3");
        let mut abort1 = testbuilder.with_abort("abort1");
        abort1.at(2).set("relay1");
        abort1.at(2).set("relay2");
        abort1.at(2).set("relay3");
        abort1.at(2).set("relay4");
        abort1
            .at(2)
            .require("sensor1", 10_f64, 20_f64, "HARD_ABORT");
        abort1
            .from(5, 50)
            .require("sensor2", 15_f64, 16_f64, "HARD_ABORT");

        let concrete_test_opt =
            ConcreteTest::try_new(testbuilder.env, testbuilder.test).to_option();
        assert!(concrete_test_opt.is_some());
    }

    #[test]
    fn build_sections() {
        let taf_output_expected = r#"abort #1:
segment +102ms
is_now_in P0, 0x38D, 0x554, #0
set P2
set P3
set P4
set P5

segment +3ms
start_check_in P1, 0x470, 0x49E, #0

segment +45ms
stop_check P1

segment +150ms
unset P2

segment +100ms
unset P3

segment +400ms
unset P4

segment +100ms
unset P5

test:
segment +102ms
is_now_in P0, 0x38D, 0x554, #0
set P2
set P3
set P4
set P5

segment +3ms
start_check_in P1, 0x470, 0x49E, #1

segment +45ms
stop_check P1

segment +150ms
unset P2

segment +100ms
unset P3

segment +400ms
unset P4

segment +100ms
unset P5

"#;

        let mut testbuilder = TestBuilder::new();
        // add some relays and sensors
        testbuilder.with_sensor("sensor1").with_sensor("sensor2");
        testbuilder
            .with_relay("relay1")
            .with_relay("relay2")
            .with_relay("relay3")
            .with_relay("relay4");

        // add sections to an abort
        let mut abort1 = testbuilder.with_abort("abort1");
        let mut abort_section = abort1.section_from(100, 1000, "relays_in_abort_section");
        abort_section.at(2).set("relay1");
        abort_section.at(2).set("relay2");
        abort_section.at(2).set("relay3");
        abort_section.at(2).set("relay4");
        abort_section
            .at(2)
            .require("sensor1", 10_f64, 20_f64, "HARD_ABORT");
        abort_section
            .from(5, 50)
            .require("sensor2", 15_f64, 16_f64, "HARD_ABORT");
        abort_section.at(200).unset("relay1");
        abort_section.at(300).unset("relay2");

        let mut nested_abort_section =
            abort_section.section_from(600, 200, "nested_section_in_abort");
        nested_abort_section.at(100).unset("relay3");
        nested_abort_section.at(200).unset("relay4");

        // Create a couple sections within the test body
        let mut test_body = testbuilder.with_body();
        let mut test_section = test_body.section_from(100, 1000, "relays_in_abort_section");
        test_section.at(2).set("relay1");
        test_section.at(2).set("relay2");
        test_section.at(2).set("relay3");
        test_section.at(2).set("relay4");
        test_section
            .at(2)
            .require("sensor1", 10_f64, 20_f64, "HARD_ABORT");
        test_section
            .from(5, 50)
            .require("sensor2", 15_f64, 16_f64, "abort1");
        test_section.at(200).unset("relay1");
        test_section.at(300).unset("relay2");

        let mut nested_test_section =
            test_section.section_from(600, 200, "nested_section_in_abort");
        nested_test_section.at(100).unset("relay3");
        nested_test_section.at(200).unset("relay4");

        let concrete_test_opt =
            ConcreteTest::try_new(testbuilder.env, testbuilder.test).to_option();
        assert!(concrete_test_opt.is_some());

        let concrete_test = concrete_test_opt.unwrap();

        let (test_body_timeline, abort_timelines) = timeline::construct_timelines(&concrete_test)
            .to_option()
            .unwrap();
        let test_opt =
            timeline::construct_test(test_body_timeline, abort_timelines, &concrete_test)
                .to_option();
        assert!(test_opt.is_some());

        let taf_output_path = NamedTempFile::new().unwrap().into_temp_path().to_path_buf();
        disassembly::emit(test_opt.unwrap(), Some(taf_output_path.clone()));

        let mut taf_output = String::new();
        fs::File::open(taf_output_path)
            .expect("Failed to open tmp file")
            .read_to_string(&mut taf_output)
            .expect("Failed to read from tmp file");

        assert_eq!(taf_output_expected, taf_output)
    }

    #[test]
    /// Test that statements are added to the correct abort stream
    fn build_multiple_aborts() {
        let taf_output_expected = r#"abort #1:
segment +10ms
set P0

segment +10ms
unset P0

abort #2:
segment +20ms
set P1

segment +10ms
unset P1

abort #3:
segment +30ms
set P2

segment +10ms
unset P2

abort #4:
segment +40ms
set P3

segment +10ms
unset P3

test:
segment +50ms
set P4

segment +10ms
unset P4

"#;

        let mut testbuilder = TestBuilder::new();
        // add some relays
        testbuilder
            .with_relay("relay1")
            .with_relay("relay2")
            .with_relay("relay3")
            .with_relay("relay4")
            .with_relay("relay5");

        let mut abort1 = testbuilder.with_abort("abort1");
        abort1.at(10).set("relay1");
        abort1.at(20).unset("relay1");

        let mut abort2 = testbuilder.with_abort("abort2");
        abort2.at(20).set("relay2");
        abort2.at(30).unset("relay2");

        let mut abort3 = testbuilder.with_abort("abort3");
        abort3.at(30).set("relay3");
        abort3.at(40).unset("relay3");

        let mut abort4 = testbuilder.with_abort("abort4");
        abort4.at(40).set("relay4");
        abort4.at(50).unset("relay4");

        let mut test_body = testbuilder.with_body();
        test_body.at(50).set("relay5");
        test_body.at(60).unset("relay5");

        let concrete_test_opt =
            ConcreteTest::try_new(testbuilder.env, testbuilder.test).to_option();
        assert!(concrete_test_opt.is_some());

        let concrete_test = concrete_test_opt.unwrap();
        let (test_body_timeline, abort_timelines) = timeline::construct_timelines(&concrete_test)
            .to_option()
            .unwrap();
        let test_opt =
            timeline::construct_test(test_body_timeline, abort_timelines, &concrete_test)
                .to_option();
        assert!(test_opt.is_some());

        let taf_output_path = NamedTempFile::new().unwrap().into_temp_path().to_path_buf();
        disassembly::emit(test_opt.unwrap(), Some(taf_output_path.clone()));

        let mut taf_output = String::new();
        fs::File::open(taf_output_path)
            .expect("Failed to open tmp file")
            .read_to_string(&mut taf_output)
            .expect("Failed to read from tmp file");

        assert_eq!(taf_output_expected, taf_output)
    }

    #[test]
    /// Test after_end()
    fn after_end_test() {
        let taf_output_expected = r#"abort #1:
segment +0ms
set P0

segment +10ms
unset P0

segment +20ms
set P0

segment +20ms
unset P0

test:
segment +50ms
is_now_in P1, 0x38D, 0x554, #0

segment +100ms
is_now_in P1, 0x38D, 0x554, #0

segment +400ms
is_now_in P1, 0x38D, 0x554, #0

"#;

        let mut testbuilder = TestBuilder::new();
        // add some relays
        testbuilder.with_relay("relay1");
        testbuilder.with_sensor("sensor1");

        let mut abort1 = testbuilder.with_abort("abort1");
        abort1.after_end(0).set("relay1");
        abort1.after_end(10).unset("relay1");
        abort1.after_end(20).set("relay1");
        abort1.after_end(20).unset("relay1");

        let mut test_body = testbuilder.with_body();
        test_body
            .after_end(50)
            .require("sensor1", 10_f64, 20_f64, "HARD_ABORT");
        test_body
            .after_end(100)
            .require("sensor1", 10_f64, 20_f64, "HARD_ABORT");
        test_body
            .after_end(400)
            .require("sensor1", 10_f64, 20_f64, "HARD_ABORT");

        let concrete_test_opt =
            ConcreteTest::try_new(testbuilder.env, testbuilder.test).to_option();
        assert!(concrete_test_opt.is_some());

        let concrete_test = concrete_test_opt.unwrap();
        let (test_body_timeline, abort_timelines) = timeline::construct_timelines(&concrete_test)
            .to_option()
            .unwrap();
        let test_opt =
            timeline::construct_test(test_body_timeline, abort_timelines, &concrete_test)
                .to_option();
        assert!(test_opt.is_some());

        let taf_output_path = NamedTempFile::new().unwrap().into_temp_path().to_path_buf();
        disassembly::emit(test_opt.unwrap(), Some(taf_output_path.clone()));

        let mut taf_output = String::new();
        fs::File::open(taf_output_path)
            .expect("Failed to open tmp file")
            .read_to_string(&mut taf_output)
            .expect("Failed to read from tmp file");

        assert_eq!(taf_output_expected, taf_output)
    }
}