//! Session Transfer Helpers
//!
//! This module contains helper functions for transferring sessions between connectors.
//! These are primarily used by the Gateway connector when selecting a target connector.

use std::time::Duration;
use tokio::sync::broadcast;

use crate::connectors::ConnectorCommand;
use dirigent_protocol::session::{SessionModeState, SessionModelState};

/// Create a session in a connector and wait for the event
///
/// Returns the session ID along with optional models/modes from the SessionCreated event.
/// The Session struct in SessionCreated already contains models/modes from the
/// connector's session/new response.
pub async fn create_session_in_connector(
    cmd_tx: &tokio::sync::mpsc::Sender<ConnectorCommand>,
    events: &mut broadcast::Receiver<dirigent_protocol::Event>,
) -> Result<(String, Option<SessionModelState>, Option<SessionModeState>), String> {
    if let Err(e) = cmd_tx
        .send(ConnectorCommand::CreateSession {
            cwd: None,
            project_id: None,
            ownership: dirigent_protocol::SessionOwnership::internal(),
        })
        .await
    {
        return Err(format!("Failed to send create command: {}", e));
    }

    // Wait for SessionCreated event
    let timeout = Duration::from_secs(30);
    let start = std::time::Instant::now();

    while start.elapsed() < timeout {
        match tokio::time::timeout(Duration::from_secs(1), events.recv()).await {
            Ok(Ok(dirigent_protocol::Event::SessionCreated { session, .. })) => {
                // Extract models/modes directly from the Session struct
                return Ok((session.id, session.models, session.modes));
            }
            Ok(Ok(dirigent_protocol::Event::Error { message })) => {
                return Err(format!("Connector error: {}", message));
            }
            Ok(Err(_)) => {
                return Err("Event channel closed".to_string());
            }
            Err(_) => continue, // Timeout on this iteration, keep waiting
            _ => continue,
        }
    }

    Err("Timeout waiting for session creation".to_string())
}

/// Wait for a session load event
///
/// Waits for a SessionUpdated or SessionCreated event for the given session ID,
/// or returns an error if a SessionError event is received or timeout occurs.
pub async fn wait_for_session_event(
    events: &mut broadcast::Receiver<dirigent_protocol::Event>,
    session_id: &str,
    timeout: Duration,
) -> Result<String, String> {
    let start = std::time::Instant::now();

    while start.elapsed() < timeout {
        match tokio::time::timeout(Duration::from_millis(500), events.recv()).await {
            Ok(Ok(dirigent_protocol::Event::SessionUpdated { session, .. }))
                if session.id == session_id =>
            {
                return Ok(session.id);
            }
            Ok(Ok(dirigent_protocol::Event::SessionCreated { session, .. }))
                if session.id == session_id =>
            {
                return Ok(session.id);
            }
            Ok(Ok(dirigent_protocol::Event::SessionError {
                session_id: sid,
                error_message,
                ..
            })) if sid == session_id => {
                return Err(error_message);
            }
            Ok(Err(_)) => return Err("Event channel closed".to_string()),
            Err(_) => continue,
            _ => continue,
        }
    }

    Err(format!("Session '{}' not found", session_id))
}

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

    #[tokio::test]
    async fn test_wait_for_session_event_timeout() {
        let (_tx, mut rx) = broadcast::channel::<dirigent_protocol::Event>(10);

        let result =
            wait_for_session_event(&mut rx, "test-session", Duration::from_millis(100)).await;

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("not found"));
    }
}
