//! The remote endpoint boundary: one trait for reaching a named target, a manager that
//! dispatches the four remote `HostRequest` variants to the registered endpoint, and the
//! relay endpoint that starts a real worker over SSH for each ensured namespace.
//!
//! Opaque call dispatch is refused until the target worker's dispatcher is present, and
//! nothing here is rerouted to local execution.

use anyhow::{bail, ensure, Context, Result};
use std::{collections::BTreeMap, time::Duration};

pub type RemoteFuture =
    std::pin::Pin<Box<dyn std::future::Future<Output = Result<serde_json::Value>> + Send>>;

pub trait RemoteEndpoint: Send + Sync {
    fn ensure(&self, session: String, workspace: Option<String>) -> RemoteFuture;
    fn call(
        &self,
        namespace_id: String,
        namespace_generation: String,
        operation_id: String,
        operation: String,
        payload: serde_json::Value,
    ) -> RemoteFuture;
    fn cancel(
        &self,
        session: String,
        execution_id: String,
        namespace_id: String,
        namespace_generation: String,
        operation_id: String,
    ) -> RemoteFuture;
    fn close(
        &self,
        session: String,
        execution_id: String,
        namespace_id: String,
        namespace_generation: String,
    ) -> RemoteFuture;

    fn shutdown(&self) -> RemoteFuture {
        Box::pin(async { Ok(serde_json::json!({"state": "closed"})) })
    }
}

/// Dispatches the four remote HostRequest variants to the registered
/// host-owned endpoint. It validates identity-shaped fields and JSON bounds,
/// but never interprets the endpoint's operation vocabulary.
type Endpoint = std::sync::Arc<dyn RemoteEndpoint>;
type Registry<K, V> = std::sync::Arc<std::sync::RwLock<BTreeMap<K, V>>>;

#[derive(Clone, Default)]
pub struct RemoteManager {
    endpoints: Registry<String, Endpoint>,
    namespaces: Registry<(String, String), (String, Endpoint)>,
}

impl RemoteManager {
    pub fn new() -> Self {
        Self::default()
    }

    /// Stops every namespace worker and clears all target/binding registries.
    pub async fn shutdown(&self) -> Result<()> {
        let endpoints = self
            .endpoints
            .read()
            .map_err(|_| anyhow::anyhow!("remote target registry lock poisoned"))?
            .values()
            .cloned()
            .collect::<Vec<_>>();
        for endpoint in endpoints {
            endpoint.shutdown().await?;
        }
        self.namespaces
            .write()
            .map_err(|_| anyhow::anyhow!("remote namespace registry lock poisoned"))?
            .clear();
        self.endpoints
            .write()
            .map_err(|_| anyhow::anyhow!("remote target registry lock poisoned"))?
            .clear();
        Ok(())
    }

    pub fn register_target(
        &self,
        target: crate::target::TargetId,
        endpoint: std::sync::Arc<dyn RemoteEndpoint>,
    ) -> Result<()> {
        self.endpoints
            .write()
            .map_err(|_| anyhow::anyhow!("remote target registry lock poisoned"))?
            .insert(target.to_string(), endpoint);
        Ok(())
    }

    /// Registers the production relay endpoint for a configured target.
    /// Callers pass the already-selected host backend; a remote placement can
    /// therefore never be silently demoted to a local adapter.
    pub fn register_relay_target(
        &self,
        target: crate::target::RemoteTarget,
        backend: std::sync::Arc<dyn crate::worker::WorkerBackend>,
        startup_timeout: Duration,
    ) -> Result<std::sync::Arc<RelayEndpoint>> {
        let endpoint =
            std::sync::Arc::new(RelayEndpoint::new(target.clone(), backend, startup_timeout));
        self.register_target(target.id, endpoint.clone())?;
        Ok(endpoint)
    }

    pub fn unregister_target(&self, target: &crate::target::TargetId) -> Result<bool> {
        Ok(self
            .endpoints
            .write()
            .map_err(|_| anyhow::anyhow!("remote target registry lock poisoned"))?
            .remove(target.as_str())
            .is_some())
    }

    fn endpoint(&self, target: &str) -> Result<std::sync::Arc<dyn RemoteEndpoint>> {
        self.endpoints
            .read()
            .map_err(|_| anyhow::anyhow!("remote target registry lock poisoned"))?
            .get(target)
            .cloned()
            .ok_or_else(|| anyhow::anyhow!("remote target {target:?} is unavailable"))
    }

    /// Handles all remote requests in one production dispatch point. The
    /// outer session/execution identities are retained by the endpoint call;
    /// an unavailable target is returned as a typed HostReply error and is
    /// never rerouted locally.
    pub async fn handle(
        &self,
        session: &crate::protocol::SessionId,
        execution_id: &crate::protocol::OperationId,
        request: crate::protocol::HostRequest,
    ) -> crate::protocol::HostReply {
        let result = self.handle_result(session, execution_id, request).await;
        match result {
            Ok(value) => crate::protocol::HostReply::Ok { value },
            Err(error) => {
                let message = format!("{error:#}");
                let code = if message.contains("remote call payload")
                    && message.contains("remote value bound")
                {
                    "payload_too_large"
                } else if message.contains("remote operation result")
                    && message.contains("remote value bound")
                {
                    "result_too_large"
                } else if message.contains("remote error") && message.contains("remote value bound")
                {
                    "error_too_large"
                } else if message.contains("unavailable") {
                    "unknown"
                } else {
                    "remote"
                };
                crate::protocol::HostReply::error(code, message)
            }
        }
    }

    async fn handle_result(
        &self,
        session: &crate::protocol::SessionId,
        execution_id: &crate::protocol::OperationId,
        request: crate::protocol::HostRequest,
    ) -> Result<serde_json::Value> {
        use crate::protocol::HostRequest;
        match request {
            HostRequest::RemoteEnsure { target, workspace } => {
                validate_target_name(&target)?;
                let endpoint = self.endpoint(&target)?;
                let value = endpoint.ensure(session.to_string(), workspace).await?;
                ensure_json_bound(&value, "remote ensure result")?;
                let object = value
                    .as_object()
                    .ok_or_else(|| anyhow::anyhow!("remote ensure result is not an object"))?;
                let namespace = object
                    .get("namespace_id")
                    .and_then(serde_json::Value::as_str)
                    .ok_or_else(|| anyhow::anyhow!("remote ensure result has no namespace_id"))?;
                let generation = object
                    .get("namespace_generation")
                    .and_then(serde_json::Value::as_str)
                    .ok_or_else(|| {
                        anyhow::anyhow!("remote ensure result has no namespace_generation")
                    })?;
                crate::target::NamespaceId::parse(namespace.to_owned())?;
                crate::target::NamespaceGeneration::parse(generation.to_owned())?;
                self.namespaces
                    .write()
                    .map_err(|_| anyhow::anyhow!("remote namespace registry lock poisoned"))?
                    .insert(
                        (namespace.to_owned(), generation.to_owned()),
                        (target, endpoint),
                    );
                Ok(value)
            }
            HostRequest::RemoteCall {
                target,
                namespace_id,
                namespace_generation,
                operation_id,
                operation,
                payload,
            } => {
                validate_target_name(&target)?;
                validate_namespace_binding(&namespace_id, &namespace_generation, &operation_id)?;
                ensure_remote_value_bound(&payload, "remote call payload")?;
                let (bound_target, endpoint) =
                    self.namespace(&namespace_id, &namespace_generation)?;
                ensure!(
                    bound_target == target,
                    "remote call target does not match namespace binding"
                );
                let value = endpoint
                    .call(
                        namespace_id,
                        namespace_generation,
                        operation_id,
                        operation,
                        payload,
                    )
                    .await?;
                ensure_json_bound(&value, "remote operation result")?;
                Ok(value)
            }
            HostRequest::RemoteCancel {
                namespace_id,
                namespace_generation,
                operation_id,
            } => {
                validate_namespace_binding(&namespace_id, &namespace_generation, &operation_id)?;
                let (_, endpoint) = self.namespace(&namespace_id, &namespace_generation)?;
                let value = endpoint
                    .cancel(
                        session.to_string(),
                        execution_id.to_string(),
                        namespace_id,
                        namespace_generation,
                        operation_id,
                    )
                    .await?;
                ensure_json_bound(&value, "remote cancellation result")?;
                Ok(value)
            }
            HostRequest::RemoteClose {
                namespace_id,
                namespace_generation,
            } => {
                validate_namespace_binding(
                    &namespace_id,
                    &namespace_generation,
                    "close-operation",
                )?;
                let key = (namespace_id.clone(), namespace_generation.clone());
                let (_, endpoint) = self.namespace(&namespace_id, &namespace_generation)?;
                let value = endpoint
                    .close(
                        session.to_string(),
                        execution_id.to_string(),
                        namespace_id,
                        namespace_generation,
                    )
                    .await?;
                ensure_json_bound(&value, "remote close result")?;
                self.namespaces
                    .write()
                    .map_err(|_| anyhow::anyhow!("remote namespace registry lock poisoned"))?
                    .remove(&key);
                Ok(value)
            }
            _ => bail!("request is not a remote namespace operation"),
        }
    }

    fn namespace(
        &self,
        namespace: &str,
        generation: &str,
    ) -> Result<(String, std::sync::Arc<dyn RemoteEndpoint>)> {
        self.namespaces
            .read()
            .map_err(|_| anyhow::anyhow!("remote namespace registry lock poisoned"))?
            .get(&(namespace.to_owned(), generation.to_owned()))
            .cloned()
            .ok_or_else(|| anyhow::anyhow!("remote namespace binding is unavailable"))
    }
}

/// A production remote namespace endpoint backed by the shared M1 worker
/// launch path. The backend is normally `RemoteBackend`, so each ensure starts
/// a real worker over SSH with the issued target binding and verified identity.
/// Opaque call dispatch is deliberately refused until the target worker's
/// dispatcher is present; this is not rerouted to local execution.
pub struct RelayEndpoint {
    target: crate::target::RemoteTarget,
    backend: std::sync::Arc<dyn crate::worker::WorkerBackend>,
    startup_timeout: Duration,
    workers: std::sync::Arc<tokio::sync::Mutex<BTreeMap<(String, String), crate::worker::Worker>>>,
}

impl RelayEndpoint {
    pub fn new(
        target: crate::target::RemoteTarget,
        backend: std::sync::Arc<dyn crate::worker::WorkerBackend>,
        startup_timeout: Duration,
    ) -> Self {
        Self {
            target,
            backend,
            startup_timeout,
            workers: std::sync::Arc::new(tokio::sync::Mutex::new(BTreeMap::new())),
        }
    }

    async fn ensure_inner(
        &self,
        session: String,
        workspace: Option<String>,
    ) -> Result<serde_json::Value> {
        let session = crate::protocol::SessionId::try_from(session)
            .map_err(|error| anyhow::anyhow!(error))?;
        let workspace = workspace.unwrap_or_else(|| self.target.workspace.path.clone());
        ensure!(
            workspace == self.target.workspace.path,
            "remote workspace does not match the configured target workspace"
        );
        let namespace_id =
            crate::target::NamespaceId::parse(format!("ns-{}", uuid::Uuid::new_v4().simple()))?;
        let namespace_generation = crate::target::NamespaceGeneration::fresh();
        let binding = crate::protocol::TargetStamp {
            target_id: self.target.id.clone(),
            namespace_id: namespace_id.clone(),
            namespace_generation: namespace_generation.clone(),
        };
        let worker = crate::worker::Worker::launch_with(
            self.backend.as_ref(),
            self.startup_timeout,
            session,
            crate::protocol::Role::Remote,
            None,
            Some(binding),
        )
        .await
        .context("launch remote relay namespace")?;
        let identity = worker
            .target_identity()
            .context("remote worker did not retain verified target identity")?;
        let capabilities = worker
            .worker_identity()
            .map(|identity| identity.runtime.capabilities)
            .unwrap_or_default();
        self.workers.lock().await.insert(
            (namespace_id.to_string(), namespace_generation.to_string()),
            worker,
        );
        Ok(serde_json::json!({
            "target": self.target.id,
            "workspace": workspace,
            "namespace_id": namespace_id,
            "namespace_generation": namespace_generation,
            "persistence": "connection-scoped",
            "identity": identity,
            "capabilities": capabilities,
        }))
    }

    async fn cancel_inner(
        &self,
        namespace_id: String,
        namespace_generation: String,
    ) -> Result<serde_json::Value> {
        let workers = self.workers.lock().await;
        ensure!(
            workers.contains_key(&(namespace_id, namespace_generation)),
            "remote namespace binding is unavailable"
        );
        // The current shared worker API has no per-operation cancellation;
        // never collapse that limitation into a false confirmation.
        Ok(serde_json::json!({
            "state": "unknown",
            "reason": "remote worker has no per-operation cancellation boundary"
        }))
    }

    async fn close_inner(
        &self,
        namespace_id: String,
        namespace_generation: String,
    ) -> Result<serde_json::Value> {
        let worker = self
            .workers
            .lock()
            .await
            .remove(&(namespace_id, namespace_generation))
            .ok_or_else(|| anyhow::anyhow!("remote namespace binding is unavailable"))?;
        worker.stop().await;
        Ok(serde_json::json!({"state": "closed"}))
    }
}

impl RemoteEndpoint for RelayEndpoint {
    fn ensure(&self, session: String, workspace: Option<String>) -> RemoteFuture {
        let this = self.clone_for_task();
        Box::pin(async move { this.ensure_inner(session, workspace).await })
    }

    fn call(
        &self,
        namespace_id: String,
        namespace_generation: String,
        operation_id: String,
        operation: String,
        payload: serde_json::Value,
    ) -> RemoteFuture {
        let workers = self.workers.clone();
        Box::pin(async move {
            ensure_remote_value_bound(&payload, "remote call payload")?;
            let operation_id = crate::protocol::OperationId::try_from(operation_id)
                .map_err(|error| anyhow::anyhow!(error))?;
            let worker = workers
                .lock()
                .await
                .get(&(namespace_id, namespace_generation))
                .cloned()
                .ok_or_else(|| anyhow::anyhow!("remote namespace binding is unavailable"))?;
            let outcome = worker
                .call(
                    crate::worker::WorkCommand::RemoteDispatch {
                        id: operation_id,
                        operation,
                        payload,
                    },
                    Duration::from_secs(60),
                    crate::worker::no_host_operations(),
                )
                .await;
            match outcome {
                Ok(crate::protocol::Outcome::Remote { value }) => {
                    ensure_remote_value_bound(&value, "remote operation result")?;
                    Ok(serde_json::json!({"status": "ok", "value": value}))
                }
                Ok(crate::protocol::Outcome::RemoteError { code, message }) => {
                    Ok(serde_json::json!({
                        "status": "error",
                        "error": {"kind": code, "message": message}
                    }))
                }
                Ok(crate::protocol::Outcome::Unknown { reason }) => Ok(serde_json::json!({
                    "status": "error",
                    "error": {"kind": "unknown", "reason": reason}
                })),
                Ok(other) => Err(anyhow::anyhow!(
                    "remote worker returned non-remote outcome: {other:?}"
                )),
                Err(error) => Ok(serde_json::json!({
                    "status": "error",
                    "error": {
                        "kind": "unknown",
                        "reason": format!("remote dispatch outcome is unknown: {error:#}")
                    }
                })),
            }
        })
    }

    fn shutdown(&self) -> RemoteFuture {
        let workers = self.workers.clone();
        Box::pin(async move {
            let workers = {
                let mut guard = workers.lock().await;
                std::mem::take(&mut *guard)
                    .into_values()
                    .collect::<Vec<_>>()
            };
            for worker in workers {
                worker.stop().await;
            }
            Ok(serde_json::json!({"state": "closed"}))
        })
    }

    fn cancel(
        &self,
        _session: String,
        _execution_id: String,
        namespace_id: String,
        namespace_generation: String,
        _operation_id: String,
    ) -> RemoteFuture {
        let this = self.clone_for_task();
        Box::pin(async move { this.cancel_inner(namespace_id, namespace_generation).await })
    }

    fn close(
        &self,
        _session: String,
        _execution_id: String,
        namespace_id: String,
        namespace_generation: String,
    ) -> RemoteFuture {
        let this = self.clone_for_task();
        Box::pin(async move { this.close_inner(namespace_id, namespace_generation).await })
    }
}

impl RelayEndpoint {
    fn clone_for_task(&self) -> Self {
        Self {
            target: self.target.clone(),
            backend: self.backend.clone(),
            startup_timeout: self.startup_timeout,
            workers: self.workers.clone(),
        }
    }
}

fn validate_target_name(target: &str) -> Result<()> {
    crate::target::TargetId::parse(target.to_owned())?;
    Ok(())
}

fn validate_namespace_binding(namespace: &str, generation: &str, operation: &str) -> Result<()> {
    crate::target::NamespaceId::parse(namespace.to_owned())?;
    crate::target::NamespaceGeneration::parse(generation.to_owned())?;
    crate::protocol::OperationId::try_from(operation.to_owned())
        .map_err(|error| anyhow::anyhow!(error))?;
    Ok(())
}

fn ensure_json_bound(value: &serde_json::Value, label: &str) -> Result<()> {
    ensure!(
        serde_json::to_vec(value)?.len() <= crate::protocol::MAX_FRAME,
        "{label} exceeds the {}-byte frame bound",
        crate::protocol::MAX_FRAME
    );
    Ok(())
}

fn ensure_remote_value_bound(value: &serde_json::Value, label: &str) -> Result<()> {
    ensure!(
        serde_json::to_vec(value)?.len() <= crate::protocol::MAX_REMOTE_VALUE,
        "{label} exceeds the {}-byte remote value bound",
        crate::protocol::MAX_REMOTE_VALUE
    );
    Ok(())
}