//! Host side of a microVM's route to the node's OWN locally-hosted model.
//!
//! # Why this is not the egress proxy
//!
//! [`crate::egress`] exists so an agent can reach the public internet under an
//! allow-list: it speaks HTTP `CONNECT`, takes a destination from the guest,
//! resolves it, and decides. Every one of those powers is a liability, which is
//! why that module is careful about ports, IP literals and suffix matching.
//!
//! This is the opposite shape. There is **no destination in the protocol**. The
//! guest opens a socket; the host connects it to `127.0.0.1:11434` on the node
//! and copies bytes. A compromised guest can ask for nothing else, because there
//! is nothing to ask — it is a pipe, not a proxy. That is strictly narrower than
//! anything the allow-list could express, and it is why routing a local model
//! through `egress` would have been the worse design: it would have meant
//! relaxing the 443-only rule and the IP-literal refusal, both of which exist
//! because a unit test caught them being bypassed.
//!
//! # Why plaintext is right here
//!
//! The bytes go guest loopback → vsock → host loopback. They never touch a
//! network, so there is no wire for TLS to protect. Ollama stays bound to
//! `127.0.0.1` on the node and is never exposed to the tailnet, which is a
//! stronger position than terminating TLS in front of it would have been.
//!
//! # Why it is per-backend
//!
//! The node binds this socket only for a backend declared to use a local model.
//! On every other backend the guest's listener is still there and simply gets a
//! refusal — the same fail-closed default `provider_hosts` applies to egress.

use std::path::{Path, PathBuf};

use tokio::net::{TcpStream, UnixListener};

/// Host-side vsock port. Must match `fcagent`'s `MODEL_VSOCK_PORT`.
pub const MODEL_PORT: u32 = 9003;

/// Where the node's model server listens. Loopback, and not configurable from
/// the guest by design — see the module docs.
const OLLAMA_ADDR: &str = "127.0.0.1:11434";

/// Whether a backend is served by a model running on the node itself.
///
/// Named individually rather than by prefix. An unrecognised backend must not
/// acquire a route to anything by accident, which is the same rule
/// `egress::provider_hosts` and `mission_runtime::microvm_credential_for`
/// already apply from their own side.
pub fn uses_local_model(backend: Option<&str>) -> bool {
    matches!(backend, Some("local-ornith"))
}

/// Bind the guest's local-model socket, if this backend has one.
///
/// `Ok(None)` means "this backend does not use a local model" and is the normal
/// case. An error means it should have had one and could not — reported by the
/// caller, never silently swallowed, because the symptom otherwise is an agent
/// that hangs on its first turn.
pub fn start(
    uds: &Path,
    vm_id: &str,
    backend: Option<&str>,
) -> Result<Option<(PathBuf, tokio::task::JoinHandle<()>)>, String> {
    if !uses_local_model(backend) {
        return Ok(None);
    }
    let path = PathBuf::from(format!("{}_{}", uds.display(), MODEL_PORT));
    // Firecracker leaves these behind exactly as it does its own socket, and a
    // stale file makes bind fail with EADDRINUSE.
    let _ = std::fs::remove_file(&path);
    let listener =
        UnixListener::bind(&path).map_err(|e| format!("bind {}: {e}", path.display()))?;

    eprintln!(
        "microvm {vm_id}: local model socket on {} -> {OLLAMA_ADDR}",
        path.display()
    );
    let vm = vm_id.to_string();
    let task = tokio::spawn(async move {
        loop {
            match listener.accept().await {
                Ok((s, _)) => {
                    let vm = vm.clone();
                    tokio::spawn(async move {
                        if let Err(e) = pipe(s).await {
                            // Loud, because the failure a mission sees is a turn
                            // that never answers. A refused connection here means
                            // the node's model server is down, and that is worth
                            // saying out loud rather than leaving to a timeout.
                            eprintln!("microvm {vm}: local model pipe failed: {e}");
                        }
                    });
                }
                Err(e) => {
                    eprintln!("microvm {vm}: local model accept failed: {e}");
                    return;
                }
            }
        }
    });
    Ok(Some((path, task)))
}

/// Splice one guest connection onto a fresh connection to the node's model.
async fn pipe(mut guest: tokio::net::UnixStream) -> Result<(), String> {
    let mut model = TcpStream::connect(OLLAMA_ADDR)
        .await
        .map_err(|e| format!("connect {OLLAMA_ADDR}: {e}"))?;
    tokio::io::copy_bidirectional(&mut guest, &mut model)
        .await
        .map(|_| ())
        .map_err(|e| format!("copy: {e}"))
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Only the backends that are meant to have a local model get one.
    ///
    /// The negative half is the point: an unrecognised backend acquiring a route
    /// to the node's own model server would be a hole opened by a typo, and it
    /// would be invisible because the mission would simply work.
    #[test]
    fn a_local_route_is_never_granted_by_accident() {
        assert!(uses_local_model(Some("local-ornith")));

        for other in [
            None,
            Some(""),
            Some("default"),
            Some("claude"),
            Some("canary-claude"),
            Some("glm"),
            Some("kimi"),
            Some("local"),
            Some("local-ornith-typo"),
            Some("ornith"),
        ] {
            assert!(
                !uses_local_model(other),
                "{other:?} must not reach the node's model server"
            );
        }
    }

    /// The guest cannot name a destination, so there is nothing to validate.
    ///
    /// This asserts the property that makes this module safe enough to skip the
    /// allow-list entirely: the upstream address is a constant. If it ever
    /// becomes a parameter, this file needs everything `egress` has.
    #[test]
    fn the_upstream_address_is_a_constant_not_an_input() {
        let src = include_str!("local_model.rs");
        // Needles are split so they do not match themselves in this file.
        assert_eq!(
            src.matches(concat!("TcpStream", "::connect(")).count(),
            1,
            "exactly one dial site, and it must use the constant"
        );
        assert!(src.contains(concat!("TcpStream", "::connect(OLLAMA_ADDR)")));
        assert!(
            OLLAMA_ADDR.starts_with("127.0.0.1:"),
            "the model server must be reached on loopback only"
        );
    }
}