vex-client 0.6.5

use anyhow::{Result, bail};
use tokio::io;
use tokio::net::TcpStream;
use uuid::Uuid;
use vex_proto::{
    AgentEntry, ClientMessage, Frame, RepoEntry, ServerMessage, ShellInfo, WorkstreamInfo,
    read_frame, send_client_message,
};

/// Typed async client for communicating with the vex daemon over TCP.
pub struct DaemonClient {
    port: u16,
}

impl DaemonClient {
    pub fn new(port: u16) -> Self {
        Self { port }
    }

    pub fn port(&self) -> u16 {
        self.port
    }

    /// Open a raw TCP connection to the daemon.
    pub async fn connect(&self) -> Result<TcpStream> {
        TcpStream::connect(("127.0.0.1", self.port))
            .await
            .map_err(|e| {
                anyhow::anyhow!(
                    "could not connect to daemon on port {}: {} (is the daemon running?)",
                    self.port,
                    e
                )
            })
    }

    /// One-shot request-response: connect, send, read one control frame.
    pub async fn request(&self, msg: &ClientMessage) -> Result<ServerMessage> {
        let stream = self.connect().await?;
        let (mut reader, mut writer) = io::split(stream);
        send_client_message(&mut writer, msg).await?;
        match read_frame(&mut reader).await? {
            Some(Frame::Control(data)) => {
                let resp: ServerMessage = serde_json::from_slice(&data)?;
                Ok(resp)
            }
            Some(Frame::Data(_)) => bail!("unexpected data frame"),
            None => bail!("server closed connection"),
        }
    }

    // ── Typed convenience methods ──────────────────────────────────

    pub async fn shell_create(
        &self,
        shell: Option<String>,
        repo: Option<String>,
        workstream: Option<String>,
    ) -> Result<Uuid> {
        let resp = self
            .request(&ClientMessage::CreateShell {
                shell,
                repo,
                workstream,
            })
            .await?;
        match resp {
            ServerMessage::ShellCreated { id } => Ok(id),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }

    pub async fn shell_list(&self) -> Result<Vec<ShellInfo>> {
        let resp = self.request(&ClientMessage::ListShells).await?;
        match resp {
            ServerMessage::Shells { shells } => Ok(shells),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }

    pub async fn shell_kill(&self, id: Uuid) -> Result<()> {
        let resp = self.request(&ClientMessage::KillShell { id }).await?;
        match resp {
            ServerMessage::ShellEnded { .. } => Ok(()),
            ServerMessage::Error { message } => bail!("{}", message),
            _ => Ok(()),
        }
    }

    pub async fn agent_list(&self) -> Result<Vec<AgentEntry>> {
        let resp = self.request(&ClientMessage::AgentList).await?;
        match resp {
            ServerMessage::AgentListResponse { agents } => Ok(agents),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }

    pub async fn agent_notifications(&self) -> Result<Vec<AgentEntry>> {
        let resp = self.request(&ClientMessage::AgentNotifications).await?;
        match resp {
            ServerMessage::AgentListResponse { agents } => Ok(agents),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }

    pub async fn agent_spawn(&self, repo: &str, workstream: Option<&str>) -> Result<Uuid> {
        let resp = self
            .request(&ClientMessage::AgentSpawn {
                repo: repo.to_string(),
                workstream: workstream.map(String::from),
            })
            .await?;
        match resp {
            ServerMessage::ShellCreated { id } => Ok(id),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }

    pub async fn agent_prompt(&self, shell_id: Uuid, text: &str) -> Result<()> {
        let resp = self
            .request(&ClientMessage::AgentPrompt {
                shell_id,
                text: text.to_string(),
            })
            .await?;
        match resp {
            ServerMessage::AgentPromptSent { .. } => Ok(()),
            ServerMessage::Error { message } => bail!("{}", message),
            _ => Ok(()),
        }
    }

    pub async fn repo_list(&self) -> Result<Vec<RepoEntry>> {
        let resp = self.request(&ClientMessage::RepoList).await?;
        match resp {
            ServerMessage::Repos { repos } => Ok(repos),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }

    pub async fn repo_add(&self, name: &str, path: &std::path::Path) -> Result<()> {
        let resp = self
            .request(&ClientMessage::RepoAdd {
                name: name.to_string(),
                path: path.to_path_buf(),
            })
            .await?;
        match resp {
            ServerMessage::RepoAdded { .. } => Ok(()),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }

    pub async fn workstream_create(&self, repo: &str, name: &str) -> Result<()> {
        let resp = self
            .request(&ClientMessage::WorkstreamCreate {
                repo: repo.to_string(),
                name: name.to_string(),
            })
            .await?;
        match resp {
            ServerMessage::WorkstreamCreated { .. } => Ok(()),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }

    pub async fn repo_remove(&self, name: &str) -> Result<()> {
        let resp = self
            .request(&ClientMessage::RepoRemove {
                name: name.to_string(),
            })
            .await?;
        match resp {
            ServerMessage::RepoRemoved { .. } => Ok(()),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }

    pub async fn workstream_remove(&self, repo: &str, name: &str) -> Result<()> {
        let resp = self
            .request(&ClientMessage::WorkstreamRemove {
                repo: repo.to_string(),
                name: name.to_string(),
            })
            .await?;
        match resp {
            ServerMessage::WorkstreamRemoved { .. } => Ok(()),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }

    pub async fn workstream_list(&self, repo: Option<&str>) -> Result<Vec<WorkstreamInfo>> {
        let resp = self
            .request(&ClientMessage::WorkstreamList {
                repo: repo.map(String::from),
            })
            .await?;
        match resp {
            ServerMessage::Workstreams { workstreams } => Ok(workstreams),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }

    pub async fn repo_introspect_path(
        &self,
        path: &std::path::Path,
    ) -> Result<(
        String,
        std::path::PathBuf,
        Option<String>,
        Option<String>,
        Vec<String>,
    )> {
        let resp = self
            .request(&ClientMessage::RepoIntrospectPath {
                path: path.to_path_buf(),
            })
            .await?;
        match resp {
            ServerMessage::RepoIntrospected {
                suggested_name,
                path,
                git_remote,
                git_branch,
                children,
            } => Ok((suggested_name, path, git_remote, git_branch, children)),
            ServerMessage::Error { message } => bail!("{}", message),
            other => bail!("unexpected response: {:?}", other),
        }
    }
}