omniference 0.1.1

use crate::adapter::ChatAdapter;
use crate::middleware::{ChatStream, RequestHandler};
use crate::types::ProviderKind;
use std::{collections::HashMap, sync::Arc};

#[derive(Clone, Default)]
pub struct AdapterRegistry {
    by_kind: HashMap<ProviderKind, Arc<dyn ChatAdapter>>,
}

impl AdapterRegistry {
    pub fn register(&mut self, adapter: Arc<dyn ChatAdapter>) {
        self.by_kind.insert(adapter.provider_kind(), adapter);
    }

    pub fn get(&self, kind: &ProviderKind) -> Option<Arc<dyn ChatAdapter>> {
        self.by_kind.get(kind).cloned()
    }

    pub fn list_kinds(&self) -> Vec<ProviderKind> {
        self.by_kind.keys().cloned().collect()
    }

    pub fn is_empty(&self) -> bool {
        self.by_kind.is_empty()
    }
}

#[derive(Clone)]
pub struct Router {
    pub registry: AdapterRegistry,
}

impl Router {
    pub fn new(registry: AdapterRegistry) -> Self {
        Self { registry }
    }

    pub async fn route_chat(
        &self,
        ir: crate::types::ChatRequestIR,
        cancel: tokio_util::sync::CancellationToken,
    ) -> anyhow::Result<impl futures_util::Stream<Item = crate::stream::StreamEvent> + Send + Unpin>
    {
        let kind = ir.model.provider.endpoint.kind.clone();
        let adapter = self
            .registry
            .get(&kind)
            .ok_or_else(|| anyhow::anyhow!("no adapter for {:?}", kind))?;

        tracing::info!(
            request_id = %ir.metadata.get("request_id").unwrap_or(&"unknown".to_string()),
            model_alias = %ir.model.alias,
            provider_kind = ?kind,
            "Routing chat request"
        );

        Ok(adapter.execute_chat(ir, cancel).await?)
    }
}

#[async_trait::async_trait]
impl RequestHandler for Router {
    async fn handle(
        &self,
        request: crate::types::ChatRequestIR,
        cancel: tokio_util::sync::CancellationToken,
    ) -> anyhow::Result<ChatStream> {
        let stream = self.route_chat(request, cancel).await?;
        Ok(Box::new(stream))
    }
}