dk-protocol 0.3.0

use std::sync::Arc;

use dk_engine::conflict::{ClaimTracker, LocalClaimTracker};
use dk_engine::repo::Engine;
use tonic::{Request, Response, Status};

use crate::auth::AuthConfig;
use crate::events::EventBus;
use crate::session::{AgentSession, SessionManager};

/// The gRPC server that implements the `AgentService` trait generated by
/// tonic from the `agent.proto` definition.
///
/// Holds a shared reference to the [`Engine`] (internally concurrent via
/// fine-grained locks) and a [`SessionManager`] for stateful agent sessions.
pub struct ProtocolServer {
    pub(crate) engine: Arc<Engine>,
    pub(crate) session_mgr: Arc<SessionManager>,
    pub(crate) auth_config: AuthConfig,
    pub(crate) event_bus: Arc<EventBus>,
    pub(crate) claim_tracker: Arc<dyn ClaimTracker>,
}

impl ProtocolServer {
    /// Create a new `ProtocolServer`.
    ///
    /// `auth_config` controls how agents authenticate on every
    /// `ConnectRequest`.  The session timeout is fixed at 30 minutes.
    pub fn new(engine: Arc<Engine>, auth_config: AuthConfig) -> Self {
        Self {
            engine,
            session_mgr: Arc::new(SessionManager::new(std::time::Duration::from_secs(
                30 * 60,
            ))),
            auth_config,
            event_bus: Arc::new(EventBus::new()),
            claim_tracker: Arc::new(LocalClaimTracker::new()),
        }
    }

    /// Borrow the engine.
    pub fn engine(&self) -> &Engine {
        &self.engine
    }

    /// Borrow the session manager.
    pub fn session_mgr(&self) -> &SessionManager {
        &self.session_mgr
    }

    /// Borrow the shared event bus.
    pub fn event_bus(&self) -> &EventBus {
        &self.event_bus
    }

    /// Return a cloned `Arc` handle to the shared event bus.
    ///
    /// Callers that need to publish to the engine event bus from a spawned
    /// task (where a borrow of `&ProtocolServer` is not available) should
    /// use this to obtain a cheaply-clonable, ownership-safe handle.
    pub fn event_bus_arc(&self) -> Arc<EventBus> {
        Arc::clone(&self.event_bus)
    }

    /// Borrow the shared symbol claim tracker.
    pub fn claim_tracker(&self) -> &dyn ClaimTracker {
        &*self.claim_tracker
    }

    /// Create a `ProtocolServer` with a custom claim tracker implementation.
    pub fn with_claim_tracker(
        engine: Arc<Engine>,
        auth_config: AuthConfig,
        claim_tracker: Arc<dyn ClaimTracker>,
    ) -> Self {
        Self {
            engine,
            session_mgr: Arc::new(SessionManager::new(std::time::Duration::from_secs(
                30 * 60,
            ))),
            auth_config,
            event_bus: Arc::new(EventBus::new()),
            claim_tracker,
        }
    }

    /// Validate an auth token against the configured secret.
    pub(crate) fn validate_auth(&self, token: &str) -> Result<String, Status> {
        self.auth_config.validate(token)
    }

    /// Check whether the bearer token in `metadata` carries admin scope.
    ///
    /// Extracts the `authorization: Bearer <token>` header, performs a full
    /// signature + expiry validation using the server's configured secret, and
    /// returns `true` when the `scope` claim equals `"admin"` or contains the
    /// word `"admin"`.
    ///
    /// Returns `false` (not an error) if there is no auth header, the token
    /// is not a JWT, or signature validation fails — callers can then fall
    /// through to the owner-check.
    pub(crate) fn has_admin_scope(&self, metadata: &tonic::metadata::MetadataMap) -> bool {
        let Some(val) = metadata.get("authorization") else {
            return false;
        };
        let Ok(header) = val.to_str() else {
            return false;
        };
        let token = header.strip_prefix("Bearer ").unwrap_or(header);
        // Perform full signature validation via the configured AuthConfig; only
        // tokens that pass cryptographic verification can claim admin scope.
        self.auth_config
            .validate_claims(token)
            .as_ref()
            .map(crate::auth::claims_have_admin_scope)
            .unwrap_or(false)
    }

    /// Look up a session by its string-encoded UUID.  Returns an error if the
    /// ID is malformed or the session has expired / does not exist.
    pub(crate) fn validate_session(&self, session_id_str: &str) -> Result<AgentSession, Status> {
        let sid = session_id_str
            .parse::<uuid::Uuid>()
            .map_err(|_| Status::invalid_argument("Invalid session ID format"))?;
        self.session_mgr
            .get_session(&sid)
            .ok_or_else(|| Status::not_found("Session not found or expired"))
    }
}

// ── AgentService tonic trait implementation ──

#[tonic::async_trait]
impl crate::agent_service_server::AgentService for ProtocolServer {
    async fn connect(
        &self,
        request: Request<crate::ConnectRequest>,
    ) -> Result<Response<crate::ConnectResponse>, Status> {
        crate::connect::handle_connect(self, request.into_inner()).await
    }

    async fn context(
        &self,
        request: Request<crate::ContextRequest>,
    ) -> Result<Response<crate::ContextResponse>, Status> {
        crate::context::handle_context(self, request.into_inner()).await
    }

    async fn submit(
        &self,
        request: Request<crate::SubmitRequest>,
    ) -> Result<Response<crate::SubmitResponse>, Status> {
        crate::submit::handle_submit(self, request.into_inner()).await
    }

    type VerifyStream = tokio_stream::wrappers::ReceiverStream<Result<crate::VerifyStepResult, Status>>;

    async fn verify(
        &self,
        request: Request<crate::VerifyRequest>,
    ) -> Result<Response<Self::VerifyStream>, Status> {
        let req = request.into_inner();
        let (tx, rx) = tokio::sync::mpsc::channel(32);

        let server_clone = ProtocolServer {
            engine: self.engine.clone(),
            session_mgr: self.session_mgr.clone(),
            auth_config: self.auth_config.clone(),
            event_bus: self.event_bus.clone(),
            claim_tracker: self.claim_tracker.clone(),
        };

        tokio::spawn(async move {
            crate::verify::handle_verify(&server_clone, req, tx).await;
        });

        Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new(rx)))
    }

    async fn merge(
        &self,
        request: Request<crate::MergeRequest>,
    ) -> Result<Response<crate::MergeResponse>, Status> {
        let resp = crate::merge::handle_merge(self, request.into_inner()).await?;
        Ok(Response::new(resp))
    }

    type WatchStream = tokio_stream::wrappers::ReceiverStream<Result<crate::WatchEvent, Status>>;

    async fn watch(
        &self,
        request: Request<crate::WatchRequest>,
    ) -> Result<Response<Self::WatchStream>, Status> {
        let req = request.into_inner();
        let (tx, rx) = tokio::sync::mpsc::channel(64);
        let server_clone = ProtocolServer {
            engine: self.engine.clone(),
            session_mgr: self.session_mgr.clone(),
            auth_config: self.auth_config.clone(),
            event_bus: self.event_bus.clone(),
            claim_tracker: self.claim_tracker.clone(),
        };
        tokio::spawn(async move {
            crate::watch::handle_watch(&server_clone, req, tx).await;
        });
        Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new(rx)))
    }

    async fn file_read(
        &self,
        request: Request<crate::FileReadRequest>,
    ) -> Result<Response<crate::FileReadResponse>, Status> {
        crate::file_read::handle_file_read(self, request.into_inner()).await
    }

    async fn file_write(
        &self,
        request: Request<crate::FileWriteRequest>,
    ) -> Result<Response<crate::FileWriteResponse>, Status> {
        crate::file_write::handle_file_write(self, request.into_inner()).await
    }

    async fn file_list(
        &self,
        request: Request<crate::FileListRequest>,
    ) -> Result<Response<crate::FileListResponse>, Status> {
        crate::file_list::handle_file_list(self, request.into_inner()).await
    }

    async fn pre_submit_check(
        &self,
        request: Request<crate::PreSubmitCheckRequest>,
    ) -> Result<Response<crate::PreSubmitCheckResponse>, Status> {
        crate::pre_submit::handle_pre_submit_check(self, request.into_inner()).await
    }

    async fn get_session_status(
        &self,
        request: Request<crate::SessionStatusRequest>,
    ) -> Result<Response<crate::SessionStatusResponse>, Status> {
        crate::session_status::handle_get_session_status(self, request.into_inner()).await
    }

    async fn push(
        &self,
        request: Request<crate::PushRequest>,
    ) -> Result<Response<crate::PushResponse>, Status> {
        let resp = crate::push::handle_push(self, request.into_inner()).await?;
        Ok(Response::new(resp))
    }

    async fn approve(
        &self,
        _request: Request<crate::ApproveRequest>,
    ) -> Result<Response<crate::ApproveResponse>, Status> {
        Err(Status::unimplemented(
            "approve is a platform-level operation; use the managed server",
        ))
    }

    async fn resolve(
        &self,
        _request: Request<crate::ResolveRequest>,
    ) -> Result<Response<crate::ResolveResponse>, Status> {
        Err(Status::unimplemented(
            "resolve is a platform-level operation; use the managed server",
        ))
    }

    async fn close(
        &self,
        _request: Request<crate::CloseRequest>,
    ) -> Result<Response<crate::CloseResponse>, Status> {
        Err(Status::unimplemented(
            "close is a platform-level operation; use the managed server",
        ))
    }

    async fn abandon(
        &self,
        request: Request<crate::AbandonRequest>,
    ) -> Result<Response<crate::AbandonResponse>, Status> {
        crate::abandon::handle_abandon(self, request).await
    }

    async fn review(
        &self,
        _request: Request<crate::ReviewRequest>,
    ) -> Result<Response<crate::ReviewResponse>, Status> {
        Err(Status::unimplemented(
            "review is a platform-level operation; use the managed server",
        ))
    }

    async fn record_review(
        &self,
        _request: Request<crate::RecordReviewRequest>,
    ) -> Result<Response<crate::RecordReviewResponse>, Status> {
        Err(Status::unimplemented(
            "record_review is a platform-level operation; use the managed server",
        ))
    }
}