little-kitty 0.0.3

use super::{
    super::{io::*, utils::*},
    control_value::*,
    controls::*,
    response::*,
};

use std::{env::*, io::*, path::*};

const MAX_PAYLOAD_CHUNK_SIZE: usize = 4096;
const MAX_PAYLOAD_CHUNK_SIZE_U64: u64 = MAX_PAYLOAD_CHUNK_SIZE as u64;

type XY = (usize, usize);

//
// Command
//

/// Kitty graphics protocol APC (Application Programming Command).
#[derive(Clone, Debug, Default, Hash, Eq, PartialEq)]
pub struct Command {
    controls: Controls,
    expect_response: bool,
    position: Option<XY>,
    tmux: Option<bool>,
}

impl Command {
    /// Add control.
    ///
    /// See [reference](https://sw.kovidgoyal.net/kitty/graphics-protocol/#control-data-reference).
    pub fn add_control<ControlT>(&mut self, key: char, value: ControlT)
    where
        ControlT: Into<ControlValue>,
    {
        self.controls.add(key, value);
    }

    /// Add control.
    ///
    /// See [reference](https://sw.kovidgoyal.net/kitty/graphics-protocol/#control-data-reference).
    ///
    /// Chainable.
    pub fn with_control<ControlT>(mut self, key: char, value: ControlT) -> Self
    where
        ControlT: Into<ControlValue>,
    {
        self.add_control(key, value);
        self
    }

    /// Whether we expect a response.
    pub fn expect_response(&self) -> bool {
        self.expect_response
    }

    /// Set whether we expect a response.
    pub fn set_expect_response(&mut self) {
        self.expect_response = true;
    }

    /// Set whether we expect a response.
    ///
    /// Chainable.
    pub fn with_expect_response(mut self) -> Self {
        self.set_expect_response();
        self
    }

    /// Position.
    pub fn position(&self) -> Option<XY> {
        self.position
    }

    /// Set the position.
    pub fn set_position(&mut self, x: usize, y: usize) {
        self.position = Some((x, y));
    }

    /// Set the position.
    ///
    /// Chainable.
    pub fn with_position(mut self, x: usize, y: usize) -> Self {
        self.set_position(x, y);
        self
    }

    /// Whether we should use [tmux](https://github.com/tmux/tmux) passthrough mode.
    ///
    /// When [None] will automatically use it if we are running in tmux.
    pub fn tmux(&self) -> Option<bool> {
        self.tmux
    }

    /// Set whether we should use [tmux](https://github.com/tmux/tmux) passthrough mode.
    ///
    /// Note that some supported terminals will accept it even when not in tmux.
    ///
    /// When [None] will automatically use it if we are running in tmux.
    pub fn set_tmux(&mut self, tmux: Option<bool>) {
        self.tmux = tmux;
    }

    /// Set whether we should use [tmux](https://github.com/tmux/tmux) passthrough mode.
    ///
    /// Note that some supported terminals will accept it even when not in tmux.
    ///
    /// When [None] will automatically use it if we are running in tmux.
    ///
    /// Chainable.
    pub fn with_tmux(mut self, tmux: Option<bool>) -> Self {
        self.set_tmux(tmux);
        self
    }

    /// Execute.
    ///
    /// If we are expecting a response then we will wait for it.
    pub fn execute(&self) -> Result<Option<Response>> {
        let session = self.session()?;

        {
            let in_tmux = self.in_tmux();
            let mut writer = self.writer(in_tmux)?;
            writer.write_start(in_tmux, self.position)?;
            self.controls.write(&mut writer)?;
            writer.write_end(in_tmux)?;
        }

        session.try_into()
    }

    /// Execute with payload.
    ///
    /// If we are expecting a response then we will wait for it.
    pub fn execute_with_payload(&self, payload: &[u8]) -> Result<Option<Response>> {
        if payload.len() > MAX_PAYLOAD_CHUNK_SIZE {
            self.execute_with_payload_from(payload)
        } else {
            let session = self.session()?;

            {
                let in_tmux = self.in_tmux();
                let mut writer = self.writer(in_tmux)?;
                writer.write_start(in_tmux, self.position)?;
                self.controls.write(&mut writer)?;
                writer.write_all(b";")?;
                writer = writer.write_base64(payload)?;
                writer.write_end(in_tmux)?;
            }

            session.try_into()
        }
    }

    /// Execute with payload from a reader.
    ///
    /// If we are expecting a response then we will wait for it.
    pub fn execute_with_payload_from<ReadT>(&self, mut payload: ReadT) -> Result<Option<Response>>
    where
        ReadT: Read,
    {
        let session = self.session()?;

        {
            let in_tmux = self.in_tmux();
            let mut writer = self.writer(in_tmux)?;

            writer.write_start(in_tmux, self.position)?;
            if !self.controls.is_empty() {
                self.controls.write(&mut writer)?;
                writer.write_all(b",")?;
            }
            writer.write_all(b"m=1;")?; // more = true

            let mut chunk = Vec::with_capacity(MAX_PAYLOAD_CHUNK_SIZE);

            loop {
                // Read next chunk
                chunk.clear();
                payload = payload.read_chunk(MAX_PAYLOAD_CHUNK_SIZE_U64, &mut chunk)?;

                if chunk.is_empty() {
                    // End chunk
                    writer.write_end(in_tmux)?;

                    // End payload
                    writer.write_start(in_tmux, None)?;
                    writer.write_all(b"m=0;")?; // more = false
                    writer.write_end(in_tmux)?;
                    break;
                }

                // End chunk
                writer = writer.write_base64(&chunk)?;
                writer.write_end(in_tmux)?;

                // Start new chunk
                writer.write_start(in_tmux, None)?;
                writer.write_all(b"m=1;")?; // more = true
            }
        }

        session.try_into()
    }

    /// Execute with path payload.
    ///
    /// If we are expecting a response then we will wait for it.
    pub fn execute_with_path_payload<PathT>(&self, path: PathT) -> Result<Option<Response>>
    where
        PathT: AsRef<Path>,
    {
        // Notes:
        // * The protocol specification doesn't mention string encoding
        // * It also doesn't mention that the path must be absolute, but implementations seem to expect it
        self.execute_with_payload(absolute(path)?.as_os_str().as_encoded_bytes())
    }

    /// Check whether the terminal supports the Kitty graphics protocol.
    ///
    /// This is done by sending a Kitty query followed by a VT100 query. Supporting terminals
    /// should respond to both queries in sequence. Other terminals should respond to only the
    /// VT100 query.
    pub fn is_supported(&self) -> Result<bool> {
        let session = KittySession::new(true)?;

        {
            let mut writer = stdout().lock();
            let in_tmux = self.in_tmux();
            writer.write_start(in_tmux, self.position)?;
            writer.write_all(b"a=q,i=1")?; // will always respond
            writer.write_end_and_device_query(in_tmux)?;
        }

        session.reader.expect("reader").has_kitty_and_device_query_responses()
    }

    fn in_tmux(&self) -> bool {
        // When in tmux the "TMUX" environment variable will bse set
        self.tmux.unwrap_or_else(|| var_os("TMUX").is_some())
    }

    fn session(&self) -> Result<KittySession> {
        KittySession::new(self.expect_response)
    }

    fn writer(&self, in_tmux: bool) -> Result<PositionalWriter<StdoutLock<'static>>> {
        let writer = stdout().lock();
        if in_tmux {
            // We will be passing through the ANSI cursor move command in write_start()
            Ok(PositionalWriter::new(writer))
        } else {
            PositionalWriter::new_save_at_maybe(writer, self.position)
        }
    }
}