//! Blocking SSH2 exec transport. No shell/process on the WASI host.
//! Host trust uses the existing lock-free known_hosts check/append policy.
use super::command::{CommandOutput, timeout};
use super::{ssh_hosts, ssh_kex, ssh_wire::*};
use std::{
    cell::RefCell,
    collections::VecDeque,
    io::{self, Read, Write},
    net::{TcpStream, ToSocketAddrs},
    rc::Rc,
    time::{Duration, Instant},
};
const WINDOW: u32 = 1024 * 1024;
include!("ssh_channel_tests.inc");

pub trait Stream: Read + Write {}
impl<T: Read + Write> Stream for T {}
/// TCP connection and per-operation I/O budgets. The overall budget starts after DNS.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TcpTimeouts {
    pub overall: Duration,
    pub connect: Duration,
    pub read: Duration,
    pub write: Duration,
}
impl Default for TcpTimeouts {
    fn default() -> Self {
        Self {
            overall: Duration::from_secs(5),
            connect: Duration::from_secs(2),
            read: Duration::from_secs(15),
            write: Duration::from_secs(15),
        }
    }
}
pub fn connect_tcp(host: &str, port: u16) -> io::Result<Box<dyn Stream>> {
    connect_tcp_with_timeouts(host, port, TcpTimeouts::default())
}
pub fn connect_tcp_with_timeouts(
    host: &str,
    port: u16,
    timeouts: TcpTimeouts,
) -> io::Result<Box<dyn Stream>> {
    Ok(Box::new(tcp_stream(host, port, timeouts)?))
}
fn tcp_stream(host: &str, port: u16, timeouts: TcpTimeouts) -> io::Result<TcpStream> {
    let addresses = (host, port)
        .to_socket_addrs()
        .map_err(|_| fail("DNS resolution failed"))?;
    let deadline = Instant::now() + timeouts.overall;
    for address in addresses.take(8) {
        let Some(remaining) = deadline.checked_duration_since(Instant::now()) else {
            break;
        };
        if let Ok(stream) = TcpStream::connect_timeout(&address, remaining.min(timeouts.connect)) {
            stream.set_read_timeout(Some(timeouts.read))?;
            stream.set_write_timeout(Some(timeouts.write))?;
            return Ok(stream);
        }
    }
    Err(fail(&format!(
        "TCP connection failed or timed out ({}s)",
        timeouts.overall.as_secs()
    )))
}
#[cfg(test)]
mod timeout_tests {
    use super::*;
    #[test]
    fn tcp_socket_uses_default_and_overridden_io_timeouts() {
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        for timeouts in [
            TcpTimeouts::default(),
            TcpTimeouts {
                overall: Duration::from_secs(7),
                connect: Duration::from_secs(3),
                read: Duration::from_secs(11),
                write: Duration::from_secs(13),
            },
        ] {
            let stream = tcp_stream("127.0.0.1", port, timeouts).unwrap();
            let (_peer, _) = listener.accept().unwrap();
            assert_eq!(stream.read_timeout().unwrap(), Some(timeouts.read));
            assert_eq!(stream.write_timeout().unwrap(), Some(timeouts.write));
        }
    }
}
pub struct Session<S> {
    pub(crate) packets: Packets<S>,
    version: Vec<u8>,
    id: Vec<u8>,
    host: String,
    port: u16,
    accept_new: bool,
    remote: u32,
    send_window: u32,
    send_max: u32,
    receive_window: u32,
    data: VecDeque<u8>,
    eof: bool,
    closed: bool,
    exit: Option<u32>,
    capture: Option<CommandOutput>,
    capture_limit: usize,
    deadline: Option<Instant>,
}
impl<S: Read + Write> Session<S> {
    pub fn start(mut stream: S, host: String, port: u16, accept_new: bool) -> io::Result<Self> {
        let version = ssh_kex::banner(&mut stream)?;
        let mut packets = Packets::new(stream);
        let mut id = Vec::new();
        ssh_kex::exchange(&mut packets, &version, &mut id, None, |key| {
            ssh_hosts::verify(&host, port, key, accept_new)
        })?;
        Ok(Self {
            packets,
            version,
            id,
            host,
            port,
            accept_new,
            remote: 0,
            send_window: 0,
            send_max: 0,
            receive_window: WINDOW,
            data: VecDeque::new(),
            eof: false,
            closed: false,
            exit: None,
            capture: None,
            capture_limit: 0,
            deadline: None,
        })
    }
    fn next(&mut self) -> io::Result<Vec<u8>> {
        for _ in 0..256 {
            if self.deadline.is_some_and(|d| Instant::now() >= d) {
                return Err(timeout());
            }
            let p = self.packets.meaningful()?;
            match p.first() {
                Some(20) => {
                    let (host, port, accept_new) = (&self.host, self.port, self.accept_new);
                    ssh_kex::exchange(
                        &mut self.packets,
                        &self.version,
                        &mut self.id,
                        Some(p),
                        |key| ssh_hosts::verify(host, port, key, accept_new),
                    )?;
                }
                Some(80) => {
                    let mut d = Decoder(&p[1..]);
                    d.string()?;
                    if d.byte()? != 0 {
                        self.packets.write(&[82])?;
                    }
                }
                Some(_) => return Ok(p),
                None => return Err(fail("empty packet")),
            }
        }
        Err(fail("too many non-progress transport messages"))
    }
    /// Authenticate without holding a private key. The callback returns an SSH-encoded signature.
    /// Async permits a vault endpoint round trip; no retries or key bytes cross this API.
    pub async fn authenticate_with<F, Fut>(
        &mut self,
        user: &str,
        public: &ssh_key::PublicKey,
        sign: F,
    ) -> io::Result<()>
    where
        F: FnOnce(Vec<u8>) -> Fut,
        Fut: std::future::Future<Output = io::Result<Vec<u8>>>,
    {
        self.packets.write(&message(5, &[b"ssh-userauth"]))?;
        let accept = self.next()?;
        if accept != message(6, &[b"ssh-userauth"]) {
            return Err(fail("userauth service rejected"));
        }
        if public.algorithm() != ssh_key::Algorithm::Ed25519 {
            return Err(fail("only ed25519 authentication supported"));
        }
        let public = public
            .to_bytes()
            .map_err(|_| fail("could not encode public key"))?;
        let mut request = message(50, &[user.as_bytes(), b"ssh-connection", b"publickey"]);
        request.push(1);
        string(&mut request, b"ssh-ed25519");
        string(&mut request, &public);
        let mut signed = Vec::new();
        string(&mut signed, &self.id);
        signed.extend_from_slice(&request);
        let signature = sign(signed).await?;
        string(&mut request, &signature);
        self.packets.write(&request)?;
        for _ in 0..16 {
            match self.next()?.first() {
                Some(52) => return Ok(()),
                Some(53) => continue, // banner, never logged
                Some(51) => return Err(fail("auth failed: publickey auth rejected")),
                _ => return Err(fail("unexpected authentication reply")),
            }
        }
        Err(fail("too many authentication banners"))
    }
    pub fn open_command(&mut self, command: &[u8]) -> io::Result<()> {
        if command.is_empty() || command.contains(&0) {
            return Err(fail("invalid exec command"));
        }
        let mut p = message(90, &[b"session"]);
        uint(&mut p, 0);
        uint(&mut p, WINDOW);
        uint(&mut p, 32768);
        self.packets.write(&p)?;
        let reply = self.next()?;
        let mut d = Decoder(&reply);
        if d.byte()? != 91 || d.uint()? != 0 {
            return Err(fail("exec channel open rejected"));
        }
        self.remote = d.uint()?;
        self.send_window = d.uint()?;
        self.send_max = d.uint()?.min(32768);
        if self.send_max == 0 {
            return Err(fail("invalid channel max packet size"));
        }
        let mut p = vec![98];
        uint(&mut p, self.remote);
        string(&mut p, b"exec");
        p.push(1);
        string(&mut p, command);
        self.packets.write(&p)?;
        for _ in 0..256 {
            let p = self.next()?;
            if p.first() == Some(&99) {
                let mut d = Decoder(&p[1..]);
                if d.uint()? != 0 {
                    return Err(fail("wrong channel in exec reply"));
                }
                return Ok(());
            }
            if p.first() == Some(&100) {
                return Err(fail("remote exec request rejected"));
            }
            self.handle(p)?;
            if self.closed {
                return Err(fail("channel closed before exec confirmation"));
            }
        }
        Err(fail("exec confirmation missing"))
    }
    /// Run exactly one exec channel, send stdin EOF, drain both streams through CLOSE.
    /// The stream must enforce the same absolute deadline for blocking reads/writes.
    pub fn execute(
        self,
        command: &[u8],
        max_bytes: usize,
        deadline: Instant,
    ) -> io::Result<CommandOutput> {
        self.execute_input(command, &[], max_bytes, deadline)
    }
    /// Execute a command with buffered stdin and capture both output streams.
    pub fn execute_input(
        mut self,
        command: &[u8],
        mut input: &[u8],
        max_bytes: usize,
        deadline: Instant,
    ) -> io::Result<CommandOutput> {
        self.deadline = Some(deadline);
        self.capture_limit = max_bytes;
        self.capture = Some(CommandOutput::default());
        self.open_command(command)?;
        while !input.is_empty() {
            let sent = self.write_data(input)?;
            input = &input[sent..];
        }
        let mut eof = vec![96];
        uint(&mut eof, self.remote);
        self.packets.write(&eof)?;
        while !self.closed {
            let p = self.next()?;
            self.handle(p)?;
        }
        let mut output = self
            .capture
            .take()
            .ok_or_else(|| fail("missing command result"))?;
        output.exit_code = self.exit;
        Ok(output)
    }
    fn replenish(&mut self, count: u32) -> io::Result<()> {
        if count == 0 || self.closed {
            return Ok(());
        }
        self.receive_window = self
            .receive_window
            .checked_add(count)
            .filter(|w| *w <= WINDOW)
            .ok_or_else(|| fail("channel receive window overflow"))?;
        let mut p = vec![93];
        uint(&mut p, self.remote);
        uint(&mut p, count);
        self.packets.write(&p)
    }
    fn handle(&mut self, p: Vec<u8>) -> io::Result<()> {
        let mut d = Decoder(&p);
        let kind = d.byte()?;
        if !(93..=100).contains(&kind) {
            return Err(fail("unexpected SSH channel message"));
        }
        if d.uint()? != 0 {
            return Err(fail("wrong SSH channel recipient"));
        }
        match kind {
            93 => {
                self.send_window = self
                    .send_window
                    .checked_add(d.uint()?)
                    .ok_or_else(|| fail("channel send window overflow"))?;
            }
            94 | 95 => {
                if self.eof {
                    return Err(fail("data after channel EOF"));
                }
                if kind == 95 {
                    d.uint()?;
                }
                let data = d.string()?;
                if data.len() > 32768 || data.len() > self.receive_window as usize {
                    return Err(fail("channel receive window exceeded"));
                }
                self.receive_window -= data.len() as u32;
                if let Some(output) = &mut self.capture {
                    let remaining = self
                        .capture_limit
                        .saturating_sub(output.stdout.len() + output.stderr.len());
                    let n = remaining.min(data.len());
                    output.truncated |= n < data.len();
                    if kind == 94 {
                        output.stdout.extend_from_slice(&data[..n]);
                    } else {
                        output.stderr.extend_from_slice(&data[..n]);
                    }
                    self.replenish(data.len() as u32)?;
                } else if kind == 94 {
                    self.data.extend(data);
                } else {
                    self.replenish(data.len() as u32)?;
                } // Git discards stderr
            }
            96 => {
                self.eof = true;
            }
            97 => {
                if !self.closed {
                    let mut p = vec![97];
                    uint(&mut p, self.remote);
                    self.packets.write(&p)?;
                }
                self.closed = true;
                self.eof = true;
            }
            98 => {
                let request = d.string()?;
                let want_reply = d.byte()? != 0;
                match request {
                    b"exit-status" => self.exit = Some(d.uint()?),
                    b"exit-signal" if self.capture.is_some() => {}
                    b"exit-signal" => return Err(fail("remote git terminated by signal")),
                    _ => {}
                }
                if want_reply {
                    let mut p = vec![100];
                    uint(&mut p, self.remote);
                    self.packets.write(&p)?;
                }
            }
            _ => return Err(fail("unsolicited channel request result")),
        }
        Ok(())
    }
    fn read_data(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }
        for _ in 0..1024 {
            if !self.data.is_empty() {
                let n = buf.len().min(self.data.len());
                for b in &mut buf[..n] {
                    *b = self
                        .data
                        .pop_front()
                        .ok_or_else(|| fail("channel buffer underflow"))?;
                }
                self.replenish(n as u32)?;
                return Ok(n);
            }
            if self.exit.is_some_and(|n| n != 0) {
                return Err(fail("remote git exited unsuccessfully"));
            }
            if self.eof {
                return Ok(0);
            }
            let p = self.next()?;
            self.handle(p)?;
        }
        Err(fail("too many non-progress channel messages"))
    }
    fn write_data(&mut self, buf: &[u8]) -> io::Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }
        if self.closed {
            return Err(fail("channel closed"));
        }
        for _ in 0..1024 {
            if self.send_window != 0 {
                let n = buf
                    .len()
                    .min(self.send_window as usize)
                    .min(self.send_max as usize);
                let mut p = vec![94];
                uint(&mut p, self.remote);
                string(&mut p, &buf[..n]);
                self.packets.write(&p)?;
                self.send_window -= n as u32;
                return Ok(n);
            }
            let p = self.next()?;
            self.handle(p)?;
            if self.closed {
                return Err(fail("channel closed"));
            }
        }
        Err(fail("channel send window stalled"))
    }
}
pub struct Reader(Rc<RefCell<Session<Box<dyn Stream>>>>);
pub struct Writer(Rc<RefCell<Session<Box<dyn Stream>>>>);
impl Read for Reader {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.borrow_mut().read_data(buf)
    }
}
impl Write for Writer {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.borrow_mut().write_data(buf)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.0.borrow_mut().packets.stream.flush()
    }
}

impl Session<Box<dyn Stream>> {
    /// Split one channel into blocking reader/writer halves, as used by gitoxide.
    pub fn split(self) -> (Reader, Writer) {
        let shared = Rc::new(RefCell::new(self));
        (Reader(shared.clone()), Writer(shared))
    }
}