sol/src/grpc/service.rs

use std::pin::Pin;
use std::sync::Arc;

use futures::Stream;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Request, Response, Status, Streaming};
use tracing::{error, info, warn};

use super::auth::Claims;
use super::proto::code_agent_server::CodeAgent;
use super::proto::*;
use super::GrpcState;

pub struct CodeAgentService {
    state: Arc<GrpcState>,
}

impl CodeAgentService {
    pub fn new(state: Arc<GrpcState>) -> Self {
        Self { state }
    }
}

#[tonic::async_trait]
impl CodeAgent for CodeAgentService {
    type SessionStream = Pin<Box<dyn Stream<Item = Result<ServerMessage, Status>> + Send>>;

    async fn session(
        &self,
        request: Request<Streaming<ClientMessage>>,
    ) -> Result<Response<Self::SessionStream>, Status> {
        // Extract JWT claims from the request extensions (set by auth middleware)
        let claims = request
            .extensions()
            .get::<Claims>()
            .cloned()
            .ok_or_else(|| Status::unauthenticated("No valid authentication token"))?;

        info!(
            user = claims.sub.as_str(),
            email = claims.email.as_deref().unwrap_or("?"),
            "New coding session"
        );

        let mut in_stream = request.into_inner();
        let state = self.state.clone();

        // Channel for sending server messages to the client
        let (tx, rx) = mpsc::channel::<Result<ServerMessage, Status>>(32);

        // Spawn the session handler
        tokio::spawn(async move {
            if let Err(e) = handle_session(&state, &claims, &mut in_stream, &tx).await {
                error!(user = claims.sub.as_str(), "Session error: {e}");
                let _ = tx
                    .send(Ok(ServerMessage {
                        payload: Some(server_message::Payload::Error(Error {
                            message: e.to_string(),
                            fatal: true,
                        })),
                    }))
                    .await;
            }
        });

        let out_stream = ReceiverStream::new(rx);
        Ok(Response::new(Box::pin(out_stream)))
    }
}

/// Handle a single coding session (runs in a spawned task).
async fn handle_session(
    state: &GrpcState,
    claims: &Claims,
    in_stream: &mut Streaming<ClientMessage>,
    tx: &mpsc::Sender<Result<ServerMessage, Status>>,
) -> anyhow::Result<()> {
    // Wait for the first message — must be StartSession
    let first = in_stream
        .message()
        .await?
        .ok_or_else(|| anyhow::anyhow!("Stream closed before StartSession"))?;

    let start = match first.payload {
        Some(client_message::Payload::Start(s)) => s,
        _ => anyhow::bail!("First message must be StartSession"),
    };

    info!(
        user = claims.sub.as_str(),
        project = start.project_path.as_str(),
        model = start.model.as_str(),
        client_tools = start.client_tools.len(),
        "Session started"
    );

    // TODO Phase 2: Create/find Matrix room for this project
    // TODO Phase 2: Create Mistral conversation
    // TODO Phase 2: Enter agent loop

    // For now, send SessionReady and echo back
    tx.send(Ok(ServerMessage {
        payload: Some(server_message::Payload::Ready(SessionReady {
            session_id: uuid::Uuid::new_v4().to_string(),
            room_id: String::new(), // TODO: Matrix room
            model: if start.model.is_empty() {
                state
                    .config
                    .agents
                    .coding_model
                    .clone()
            } else {
                start.model.clone()
            },
        })),
    }))
    .await?;

    // Main message loop
    while let Some(msg) = in_stream.message().await? {
        match msg.payload {
            Some(client_message::Payload::Input(input)) => {
                info!(
                    user = claims.sub.as_str(),
                    text_len = input.text.len(),
                    "User input received"
                );

                // TODO Phase 2: Send to Mistral, handle tool calls, stream response
                // For now, echo back as a simple acknowledgment
                tx.send(Ok(ServerMessage {
                    payload: Some(server_message::Payload::Done(TextDone {
                        full_text: format!("[stub] received: {}", input.text),
                        input_tokens: 0,
                        output_tokens: 0,
                    })),
                }))
                .await?;
            }
            Some(client_message::Payload::ToolResult(result)) => {
                info!(
                    call_id = result.call_id.as_str(),
                    is_error = result.is_error,
                    "Tool result received"
                );
                // TODO Phase 2: Feed back to Mistral
            }
            Some(client_message::Payload::Approval(approval)) => {
                info!(
                    call_id = approval.call_id.as_str(),
                    approved = approval.approved,
                    "Tool approval received"
                );
                // TODO Phase 2: Execute or skip tool
            }
            Some(client_message::Payload::End(_)) => {
                info!(user = claims.sub.as_str(), "Session ended by client");
                tx.send(Ok(ServerMessage {
                    payload: Some(server_message::Payload::End(SessionEnd {
                        summary: "Session ended.".into(),
                    })),
                }))
                .await?;
                break;
            }
            Some(client_message::Payload::Start(_)) => {
                warn!("Received duplicate StartSession — ignoring");
            }
            None => continue,
        }
    }

    Ok(())
}