//! Unix PTY ownership for interactive agent terminal sessions.

use std::os::fd::{AsFd, OwnedFd};
use std::path::Path;
use std::pin::Pin;
use std::task::{Context as TaskContext, Poll};

use anyhow::{Context, Result};
use pty_process::{Command, OwnedReadPty, OwnedWritePty, Size};
use tokio::io::AsyncWrite;
use tokio::process::Child;

/// An interactive child process and the independently owned sides of its PTY.
///
/// Keeping the write half available is important: `OwnedWritePty` also owns the
/// resize operation used by terminal WebSocket control frames.
pub struct TerminalPty {
    reader: OwnedReadPty,
    writer: TerminalWriter,
    child: Child,
}

/// Writable PTY half with the control operations needed by the agent.
pub struct TerminalWriter {
    io: OwnedWritePty,
    control: OwnedFd,
}

impl TerminalPty {
    /// Starts a configured program attached to a newly allocated PTY.
    ///
    /// Wakey retains control of the PTY streams, terminal environment, and
    /// kill-on-drop behavior; callers may only select arguments and a working
    /// directory.
    pub fn spawn(
        program: &Path,
        args: &[String],
        current_dir: Option<&Path>,
        rows: u16,
        cols: u16,
    ) -> Result<Self> {
        let (pty, pts) = pty_process::open().context("failed to open PTY")?;
        pty.resize(Size::new(rows, cols))
            .context("failed to set initial PTY size")?;
        let control = pty
            .as_fd()
            .try_clone_to_owned()
            .context("failed to duplicate PTY control descriptor")?;

        // The remote frontend is xterm.js regardless of the daemon's own
        // environment, so advertise the terminal the child actually receives.
        let mut command = Command::new(program)
            .args(args)
            .env("TERM", "xterm-256color")
            .env("COLORTERM", "truecolor")
            .kill_on_drop(true);
        if let Some(current_dir) = current_dir {
            command = command.current_dir(current_dir);
        }
        let child = command
            .spawn(pts)
            .with_context(|| format!("failed to spawn {} in PTY", program.display()))?;
        let (reader, io) = pty.into_split();

        Ok(Self {
            reader,
            writer: TerminalWriter { io, control },
            child,
        })
    }

    /// Consumes the terminal into independently driven async parts.
    pub fn into_parts(self) -> (OwnedReadPty, TerminalWriter, Child) {
        (self.reader, self.writer, self.child)
    }
}

impl TerminalWriter {
    pub fn resize(&self, rows: u16, cols: u16) -> Result<()> {
        self.io
            .resize(Size::new(rows, cols))
            .context("failed to resize PTY")
    }

    /// Requests a full-screen application to redraw after reattachment.
    ///
    /// Job-control shells give foreground programs their own process groups,
    /// so signaling the original shell group would miss programs like `btop`.
    pub fn refresh(&self) -> Result<()> {
        let foreground = nix::unistd::tcgetpgrp(&self.control)
            .context("failed to get PTY foreground process group")?;
        nix::sys::signal::killpg(foreground, nix::sys::signal::Signal::SIGWINCH)
            .context("failed to signal PTY foreground process group")
    }
}

impl AsyncWrite for TerminalWriter {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut TaskContext<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        Pin::new(&mut self.io).poll_write(cx, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<std::io::Result<()>> {
        Pin::new(&mut self.io).poll_flush(cx)
    }

    fn poll_shutdown(
        mut self: Pin<&mut Self>,
        cx: &mut TaskContext<'_>,
    ) -> Poll<std::io::Result<()>> {
        Pin::new(&mut self.io).poll_shutdown(cx)
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    use super::*;

    /// This test intentionally uses the same PTY implementation on the router.
    /// Run it remotely with:
    ///
    /// `./scripts/test_remote.ps1 -Package wakey-linux -Filter terminal::tests::pty_round_trip_and_resize -Exact -ShowOutput`
    #[tokio::test]
    #[ignore = "live system check, requires PTY and shell"]
    async fn pty_round_trip_and_resize() {
        let (mut reader, mut writer, mut child) =
            TerminalPty::spawn(Path::new("/bin/sh"), &[], None, 24, 80)
                .expect("spawn PTY shell")
                .into_parts();

        writer.resize(31, 101).expect("resize owned PTY writer");
        // Print the size from the parent shell after `stty` exits, so observing
        // the marker also means foreground control has returned to the shell.
        // Then wait for the literal probe; WINCH can interrupt a plain `read`.
        writer
            .write_all(
                b"trap 'printf \"WAKEY_WINCH\\n\"' WINCH; \
                  printf 'WAKEY_ENV:%s:%s\\n' \"$TERM\" \"$COLORTERM\"; \
                  size=$(stty size); printf 'WAKEY_SIZE:%s\\n' \"$size\"; \
                  line=; while [ \"$line\" != probe ]; do read line || line=; done; \
                  printf 'WAKEY_PTY_OK:%s\\n' \"$line\"; exit 0\n",
            )
            .await
            .expect("write shell input");
        let mut output = Vec::new();
        read_until(
            &mut reader,
            &mut output,
            "WAKEY_ENV:xterm-256color:truecolor",
        )
        .await;
        read_until(&mut reader, &mut output, "WAKEY_SIZE:31 101").await;
        writer.refresh().expect("signal foreground process group");
        writer
            .write_all(b"probe\n")
            .await
            .expect("write shell probe");
        let output = tokio::time::timeout(Duration::from_secs(5), async {
            read_until(&mut reader, &mut output, "WAKEY_PTY_OK:probe").await;
            output
        })
        .await
        .expect("PTY output timed out");
        let status = tokio::time::timeout(Duration::from_secs(5), child.wait())
            .await
            .expect("PTY child wait timed out")
            .expect("wait for PTY child");

        assert!(status.success(), "shell exited with {status}");
        let output = String::from_utf8_lossy(&output);
        assert!(
            output.contains("WAKEY_PTY_OK:probe"),
            "unexpected PTY output: {output:?}"
        );
        assert!(
            output.contains("WAKEY_SIZE:31 101"),
            "PTY did not report resized dimensions: {output:?}"
        );
        assert!(
            output.contains("WAKEY_WINCH"),
            "foreground process did not receive SIGWINCH: {output:?}"
        );
    }

    async fn read_until(reader: &mut OwnedReadPty, output: &mut Vec<u8>, marker: &str) {
        tokio::time::timeout(Duration::from_secs(5), async {
            let mut chunk = [0_u8; 4096];
            while !String::from_utf8_lossy(output).contains(marker) {
                // Use a fresh ReadBuf for each PTY read. pty-process 0.5.3's
                // AsyncRead implementation cannot safely extend Tokio's
                // partially filled buffer used by `read_to_end`.
                let count = reader.read(&mut chunk).await.expect("read PTY output");
                assert_ne!(count, 0, "PTY closed before emitting {marker}");
                output.extend_from_slice(&chunk[..count]);
            }
        })
        .await
        .unwrap_or_else(|_| panic!("PTY output timed out waiting for {marker}"));
    }
}