use anyhow::{Context, Result};
use nix::fcntl::{fcntl, FcntlArg, OFlag};
use nix::pty::{openpty, Winsize};
use nix::unistd::dup;
use std::ffi::CStr;
use std::fs::{self, File};
use std::io::{self, Read, Write};
use std::os::unix::fs::symlink;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::task::{Context as TaskContext, Poll};
use tokio::fs::remove_file;
use tokio::io::unix::AsyncFd;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tracing::info;

pub struct PtyEndpoint {
    pub async_io: AsyncPty,
    pub link_path: PathBuf,
}

pub struct AsyncPty {
    inner: AsyncFd<File>,
}

impl AsyncPty {
    pub fn new(master_fd: std::os::unix::io::RawFd) -> Result<Self> {
        fcntl(master_fd, FcntlArg::F_SETFL(OFlag::O_NONBLOCK))
            .context("failed to set PTY master non-blocking")?;

        let file = unsafe { File::from_raw_fd(master_fd) };
        let inner = AsyncFd::new(file).context("failed to wrap PTY master in AsyncFd")?;
        Ok(Self { inner })
    }
}

impl AsyncRead for AsyncPty {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut TaskContext<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        loop {
            let mut guard = match self.inner.poll_read_ready(cx) {
                Poll::Ready(Ok(guard)) => guard,
                Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                Poll::Pending => return Poll::Pending,
            };

            let unfilled = buf.initialize_unfilled();
            match guard.try_io(|inner| inner.get_ref().read(unfilled)) {
                Ok(Ok(0)) => continue,
                Ok(Ok(n)) => {
                    buf.advance(n);
                    return Poll::Ready(Ok(()));
                }
                Ok(Err(e)) => return Poll::Ready(Err(e)),
                Err(_) => continue,
            }
        }
    }
}

impl AsyncWrite for AsyncPty {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut TaskContext<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        loop {
            let mut guard = match self.inner.poll_write_ready(cx) {
                Poll::Ready(Ok(guard)) => guard,
                Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                Poll::Pending => return Poll::Pending,
            };

            match guard.try_io(|inner| inner.get_ref().write(buf)) {
                Ok(result) => return Poll::Ready(result),
                Err(_) => continue,
            }
        }
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }
}

pub fn create_pty_endpoint(link_path: &Path) -> Result<PtyEndpoint> {
    if link_path.exists() {
        anyhow::bail!(
            "link path {} already exists; remove it or choose another path",
            link_path.display()
        );
    }

    let winsize = Winsize {
        ws_row: 24,
        ws_col: 80,
        ws_xpixel: 0,
        ws_ypixel: 0,
    };

    let pty = openpty(Some(&winsize), None).context("failed to open PTY pair")?;

    let slave_path = slave_path_for_master(pty.master.as_raw_fd())?;

    if let Some(parent) = link_path.parent() {
        fs::create_dir_all(parent)
            .with_context(|| format!("failed to create parent dir {}", parent.display()))?;
    }

    symlink(&slave_path, link_path).with_context(|| {
        format!(
            "failed to symlink {} -> {}",
            link_path.display(),
            slave_path
        )
    })?;

    info!(
        link = %link_path.display(),
        slave = %slave_path,
        "PTY ready"
    );

    let owned_fd = dup(pty.master.as_raw_fd()).context("failed to dup PTY master fd")?;

    Ok(PtyEndpoint {
        async_io: AsyncPty::new(owned_fd)?,
        link_path: link_path.to_path_buf(),
    })
}

pub async fn cleanup_link(link_path: &Path) -> Result<()> {
    if link_path.exists() || link_path.is_symlink() {
        remove_file(link_path)
            .await
            .with_context(|| format!("failed to remove {}", link_path.display()))?;
        info!(link = %link_path.display(), "removed PTY symlink");
    }
    Ok(())
}

fn slave_path_for_master(master_fd: std::os::unix::io::RawFd) -> Result<String> {
    let name_ptr = unsafe { libc::ptsname(master_fd) };
    if name_ptr.is_null() {
        return Err(anyhow::anyhow!("ptsname failed for PTY master"));
    }
    let name = unsafe { CStr::from_ptr(name_ptr) };
    Ok(name.to_string_lossy().into_owned())
}
