use std::sync::OnceLock;

use nanoid::nanoid;
use paddler_types::inference_client::Message as InferenceMessage;
use paddler_types::inference_server::Message as InferenceServerMessage;
use paddler_types::inference_server::Request as InferenceServerRequest;
use paddler_types::jsonrpc::RequestEnvelope;
use paddler_types::request_params::ContinueFromConversationHistoryParams;
use paddler_types::request_params::ContinueFromRawPromptParams;
use paddler_types::request_params::GenerateEmbeddingBatchParams;
use paddler_types::request_params::continue_from_conversation_history_params::tool::tool_params::function_call::parameters_schema::validated_parameters_schema::ValidatedParametersSchema;
use reqwest::Client;
use tokio_stream::wrappers::UnboundedReceiverStream;
use url::Url;

use crate::Result;
use crate::format_api_url::format_api_url;
use crate::inference_message_stream::InferenceMessageStream;
use crate::inference_socket::pool::Pool;
use crate::stream::ndjson::Ndjson;

pub struct ClientInference<'client> {
    url: &'client Url,
    http_client: &'client Client,
    inference_socket_pool: OnceLock<Pool>,
    inference_socket_pool_size: usize,
}

impl<'client> ClientInference<'client> {
    #[must_use]
    pub const fn new(
        url: &'client Url,
        http_client: &'client Client,
        inference_socket_pool_size: usize,
    ) -> Self {
        Self {
            url,
            http_client,
            inference_socket_pool: OnceLock::new(),
            inference_socket_pool_size,
        }
    }

    pub async fn get_health(&self) -> Result<String> {
        let response = self
            .http_client
            .get(format_api_url(self.url, "/health")?)
            .send()
            .await?
            .error_for_status()?;

        Ok(response.text().await?)
    }

    fn get_inference_socket_pool(&self) -> &Pool {
        self.inference_socket_pool
            .get_or_init(|| Pool::new(self.url.clone(), self.inference_socket_pool_size))
    }

    pub async fn continue_from_conversation_history(
        &self,
        params: ContinueFromConversationHistoryParams<ValidatedParametersSchema>,
    ) -> Result<InferenceMessageStream> {
        let request_id = nanoid!();
        let message: InferenceServerMessage<ValidatedParametersSchema> =
            InferenceServerMessage::Request(RequestEnvelope {
                id: request_id.clone(),
                request: InferenceServerRequest::ContinueFromConversationHistory(params),
            });
        let rx = self
            .get_inference_socket_pool()
            .send_request(request_id, message)
            .await?;

        Ok(Box::pin(UnboundedReceiverStream::new(rx)))
    }

    pub async fn continue_from_raw_prompt(
        &self,
        params: ContinueFromRawPromptParams,
    ) -> Result<InferenceMessageStream> {
        let request_id = nanoid!();
        let message: InferenceServerMessage<ValidatedParametersSchema> =
            InferenceServerMessage::Request(RequestEnvelope {
                id: request_id.clone(),
                request: InferenceServerRequest::ContinueFromRawPrompt(params),
            });
        let rx = self
            .get_inference_socket_pool()
            .send_request(request_id, message)
            .await?;

        Ok(Box::pin(UnboundedReceiverStream::new(rx)))
    }

    pub async fn post_continue_from_conversation_history(
        &self,
        params: &ContinueFromConversationHistoryParams<ValidatedParametersSchema>,
    ) -> Result<InferenceMessageStream> {
        let response = self
            .http_client
            .post(format_api_url(
                self.url,
                "/api/v1/continue_from_conversation_history",
            )?)
            .json(params)
            .send()
            .await?
            .error_for_status()?;

        let stream = Ndjson::<InferenceMessage>::from_response(response);

        Ok(Box::pin(stream))
    }

    pub async fn post_continue_from_raw_prompt(
        &self,
        params: &ContinueFromRawPromptParams,
    ) -> Result<InferenceMessageStream> {
        let response = self
            .http_client
            .post(format_api_url(
                self.url,
                "/api/v1/continue_from_raw_prompt",
            )?)
            .json(params)
            .send()
            .await?
            .error_for_status()?;

        let stream = Ndjson::<InferenceMessage>::from_response(response);

        Ok(Box::pin(stream))
    }

    pub async fn generate_embedding_batch(
        &self,
        params: &GenerateEmbeddingBatchParams,
    ) -> Result<InferenceMessageStream> {
        let response = self
            .http_client
            .post(format_api_url(
                self.url,
                "/api/v1/generate_embedding_batch",
            )?)
            .json(params)
            .send()
            .await?
            .error_for_status()?;

        let stream = Ndjson::<InferenceMessage>::from_response(response);

        Ok(Box::pin(stream))
    }
}
