use std::io::prelude::*;
use std::io::BufReader;
use std::io::BufWriter;
use std::net::TcpStream;

use log::*;

pub struct DS1000 {
    reader: BufReader<TcpStream>,
    writer: BufWriter<TcpStream>,
}

#[derive(PartialEq, Debug)]
pub enum MDepth {
    Auto,
    Manual(usize),
}

impl DS1000 {
    pub fn new(host: &str) -> DS1000 {
        let stream = TcpStream::connect(host).unwrap();

        let reader = BufReader::new(stream.try_clone().unwrap());
        let writer = BufWriter::new(stream.try_clone().unwrap());

        DS1000 { reader, writer }
    }

    pub fn send_command(&mut self, command: &str) -> Result<(), std::io::Error> {
        trace!("Sending command: {}", command);
        self.writer.write(format!("{}\n", command).as_bytes())?;
        self.writer.flush()?;

        std::thread::sleep(std::time::Duration::from_millis(10));

        Ok(())
    }

    pub fn read_string(&mut self) -> Result<String, std::io::Error> {
        let mut line = String::new();
        self.reader.read_line(&mut line)?;
        let line = line.trim().to_string();
        trace!("Received string: {}", line);

        Ok(line)
    }

    pub fn idn(&mut self) -> std::io::Result<String> {
        self.send_command("*IDn?")?;
        self.read_string()
    }

    pub fn run(&mut self) -> std::io::Result<()> {
        self.send_command(":RUN")
    }

    pub fn stop(&mut self) -> std::io::Result<()> {
        self.send_command(":STOP")
    }

    pub fn single(&mut self) -> std::io::Result<()> {
        self.send_command(":SINGle")
    }

    pub fn clear(&mut self) -> std::io::Result<()> {
        self.send_command(":CLEar")
    }

    pub fn trigger_force(&mut self) -> std::io::Result<()> {
        self.send_command(":TFORce")
    }

    pub fn autoscale(&mut self) -> std::io::Result<()> {
        self.send_command(":AUToscale")
    }

    pub fn beep(&mut self, beep: bool) -> std::io::Result<()> {
        self.send_command(&format!(":SYSTem:BEEPer {}", if beep { 1 } else { 0 }))
    }

    pub fn get_srate(&mut self) -> std::io::Result<f64> {
        self.send_command(":ACQuire:SRATe?")?;

        let res = self.read_string()?;
        let res = res.trim();

        Ok(res.parse::<f64>().unwrap())
    }

    pub fn get_mdepth(&mut self) -> std::io::Result<MDepth> {
        self.send_command(":ACQuire:MDEPth?")?;

        let res = self.read_string()?;
        let res = res.trim();

        if res == "AUTO" {
            Ok(MDepth::Auto)
        } else {
            Ok(MDepth::Manual(res.parse::<usize>().unwrap()))
        }
    }

    pub fn set_mdepth(&mut self, mdepth: MDepth) -> std::io::Result<()> {
        match mdepth {
            MDepth::Auto => self.send_command(":ACQuire:MDEPth AUTO"),
            MDepth::Manual(mdep) => self.send_command(&format!(":ACQuire:MDEPth {}", mdep)),
        }
    }

    pub fn get_waveform(&mut self, channel: &str) -> std::io::Result<Vec<u8>> {
        self.send_command(&format!(":WAV:SOUR {}", channel))?;
        self.send_command(":WAV:MODE RAW")?;
        self.send_command(":WAV:FORM BYTE")?;

        if let MDepth::Manual(depth) = self.get_mdepth().unwrap() {
            let mut ret: Vec<u8> = Vec::with_capacity(depth);

            let n_per_iter = 250000; //from the rigol manual

            let niter = (depth as f64 / n_per_iter as f64).ceil() as usize;

            info!("Expecting {} points, {} iterations", depth, niter);

            for i in 0..niter {
                self.send_command(&format!(":WAV:STAR {}", i * n_per_iter + 1))?;
                self.send_command(&format!(":WAV:STOP {}", depth.min((i + 1) * n_per_iter)))?;
                self.send_command(":WAV:DATA?")?;

                let mut header = [0; 11];
                self.reader.read_exact(&mut header)?;

                let header_str = std::str::from_utf8(&(header[2..])).unwrap();
                trace!("header: {:?}", header_str);

                let npoints = header_str.parse::<usize>().unwrap();

                info!("Iter {}/{}, number of points: {}", i + 1, niter, npoints);

                let mut buf: Vec<u8> = Vec::with_capacity(npoints);
                buf.resize(npoints, 0);

                self.reader.read_exact(&mut buf)?;

                self.reader.read_line(&mut String::new()).unwrap();

                ret.append(&mut buf);
            }

            assert!(ret.len() == depth);
            Ok(ret)
        } else {
            panic!("Can't get waveform with auto memory depth"); //idk there's probably a way to do it
        }
    }
}

#[test]
fn test_1() {
    let mut scope = DS1000::new(&std::env::var("SCOPE_HOST").unwrap());

    scope.run().unwrap();

    scope.set_mdepth(MDepth::Manual(24000000)).unwrap();

    assert!(scope.get_mdepth().unwrap() == MDepth::Manual(24000000));

    scope.set_mdepth(MDepth::Manual(12000000)).unwrap();

    assert!(scope.get_mdepth().unwrap() == MDepth::Manual(12000000));

    scope.set_mdepth(MDepth::Auto).unwrap();

    assert!(scope.get_mdepth().unwrap() == MDepth::Auto);
}