use super::request::*;
use super::sessions::Session;
use crate::gemini::error::GeminiResponseStreamError;
use bytes::Bytes;
use derive_new::new;
use futures::Stream;
use reqwest::Response;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::{
    pin::Pin,
    task::{Context, Poll},
};

#[derive(Serialize, Deserialize, Clone, Debug)]
#[allow(non_camel_case_types)]
pub enum FinishReason {
    /// Default value. This value is unused.
    FINISH_REASON_UNSPECIFIED,
    /// Natural stop point of the model or provided stop sequence.
    STOP,
    /// The maximum number of tokens as specified in the request was reached.
    MAX_TOKENS,
    /// The response candidate content was flagged for safety reasons.
    SAFETY,
    /// The response candidate content was flagged for recitation reasons.
    RECITATION,
    /// The response candidate content was flagged for using an unsupported language.
    LANGUAGE,
    /// Unknown reason.
    OTHER,
    /// Token generation stopped because the content contains forbidden terms.
    BLOCKLIST,
    /// Token generation stopped for potentially containing prohibited content.
    PROHIBITED_CONTENT,
    /// Token generation stopped because the content potentially contains Sensitive Personally Identifiable Information (SPII).
    SPII,
    /// The function call generated by the model is invalid.
    MALFORMED_FUNCTION_CALL,
    /// Token generation stopped because generated images contain safety violations.
    IMAGE_SAFETY,
}

#[derive(Serialize, Deserialize, Clone, Debug, new)]
#[allow(non_snake_case)]
struct Candidate {
    content: Chat,
    pub finishReason: Option<FinishReason>,
}

#[derive(Serialize, Deserialize, Clone, Debug)]
#[allow(non_snake_case)]
pub struct GeminiResponse {
    candidates: Vec<Candidate>,
    pub usageMetadata: Value,
    pub modelVersion: String,
    pub promptFeedback: Option<Value>,
}
impl GeminiResponse {
    pub(crate) async fn new(response: Response) -> Result<GeminiResponse, reqwest::Error> {
        response.json().await
    }
    pub(crate) fn from_str(string: impl AsRef<str>) -> Result<Self, serde_json::Error> {
        serde_json::from_str(string.as_ref())
    }
    ///Instead use get_chat()
    #[deprecated]
    pub fn get_parts(&self) -> &Vec<Part> {
        self.candidates[0].content.parts()
    }
    pub fn get_chat(&self) -> &Chat {
        &self.candidates[0].content
    }
    pub fn get_json<T>(&self) -> Result<T, serde_json::Error>
    where
        T: serde::de::DeserializeOwned,
    {
        Self::parse_json(self.get_chat().parts())
    }
    pub fn parse_json<T>(parts: &[Part]) -> Result<T, serde_json::Error>
    where
        T: serde::de::DeserializeOwned,
    {
        let unescaped_str = Chat::extract_text_all(parts, "")
            .replace("\\\"", "\"")
            .replace("\\n", "\n");
        serde_json::from_str::<T>(&unescaped_str)
    }
    /// Instead use
    /// ```ignore
    /// geminiResponse.get_chat().is_thinking()
    /// ```
    #[deprecated]
    pub fn is_thinking(&self) -> bool {
        let parts = self.get_chat().parts();

        // If there are no text parts in the chunk at all, it is not a "thought".
        if !parts.iter().any(|p| matches!(p, Part::text(_))) {
            return false;
        }

        // If every text part in this chunk is a "thought",
        // then we consider the entire chunk a "thought".
        parts.iter().all(|part| match part {
            Part::text(text_part) => *text_part.thought(),
            _ => true,
        })
    }
    /// Instead use
    /// ```ignore
    /// geminiResponse.get_chat().get_text_no_think()
    /// ```
    ///`seperator` used to concatenate all text parts. TL;DR use "" as seperator.
    #[deprecated]
    pub fn extract_text_no_think(parts: &[Part], seperator: impl AsRef<str>) -> String {
        parts
            .iter()
            .filter_map(|part| {
                // Looking only for text parts
                if let Part::text(text_part) = part {
                    if !text_part.thought() {
                        Some(text_part.text().as_str())
                    } else {
                        None // Ignoring "thought" parts
                    }
                } else {
                    None
                }
            })
            .collect::<Vec<&str>>()
            .join(seperator.as_ref())
    }
    /// Instead use
    /// ```ignore
    /// chat.get_text_all()
    /// ```
    ///`seperator` used to concatenate all text parts. TL;DR use "" as seperator.
    #[deprecated]
    pub fn extract_text(parts: &[Part], seperator: impl AsRef<str>) -> String {
        parts
            .iter()
            .filter_map(|part| {
                if let Part::text(text_part) = part {
                    // Just return the text, without checking the `thought` flag
                    Some(text_part.text().as_str())
                } else {
                    None
                }
            })
            .collect::<Vec<&str>>()
            .join(seperator.as_ref())
    }
    /// Instead use
    /// ```ignore
    /// geminiResponse.get_chat().get_text_all()
    /// ```
    ///`seperator` used to concatenate all text parts. TL;DR use "" as seperator.
    #[deprecated]
    pub fn get_text(&self, seperator: impl AsRef<str>) -> String {
        self.get_chat().get_text_all(seperator)
    }

    /// Instead use
    /// ```ignore
    /// geminiResponse.get_chat().get_text_no_think()
    /// ```
    /// Don't include thoughts
    ///`seperator` used to concatenate all text parts. TL;DR use "" as seperator.
    #[deprecated]
    pub fn get_text_no_think(&self, seperator: impl AsRef<str>) -> String {
        self.get_chat().get_text_no_think(seperator)
    }
}

pin_project_lite::pin_project! {
    pub struct ResponseStream<F,T>
        where F:FnMut(&Session, GeminiResponse) -> T{
        #[pin]
        response_stream:Box<dyn Stream<Item = Result<Bytes, reqwest::Error>> + Unpin + Send + 'static>,
        session: Session,
        data_extractor: F,
        buffer: Vec<u8>,
    }
}
impl<F, T> Stream for ResponseStream<F, T>
where
    F: FnMut(&Session, GeminiResponse) -> T,
{
    type Item = Result<T, GeminiResponseStreamError>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut this = self.project();

        loop {
            // Look for the delimiter `\r\n\r\n` in the buffer.
            if let Some(end_of_message_pos) = this
                .buffer
                .windows(4)
                .position(|window| window == b"\r\n\r\n")
            {
                let message_bytes = this
                    .buffer
                    .drain(..end_of_message_pos + 4)
                    .collect::<Vec<u8>>();

                // Convert to string and remove the "data: " prefix and extra spaces.
                if let Some(json_str) = String::from_utf8_lossy(&message_bytes)
                    .strip_prefix("data:")
                    .map(|s| s.trim())
                {
                    if !json_str.is_empty() {
                        // Parse JSON.
                        let response = match GeminiResponse::from_str(json_str) {
                            Ok(resp) => resp,
                            Err(e) => {
                                let err = GeminiResponseStreamError::InvalidResposeFormat(format!(
                                    "JSON parsing error [{}]: {}",
                                    e, json_str
                                ));
                                return Poll::Ready(Some(Err(err)));
                            }
                        };

                        // Update the session and return the data.
                        this.session.update(&response);
                        let data = (this.data_extractor)(this.session, response);
                        return Poll::Ready(Some(Ok(data)));
                    }
                }
                continue;
            }

            // If the complete message is not in the buffer, read data from the network.
            match this.response_stream.as_mut().poll_next(cx) {
                Poll::Ready(Some(Ok(bytes))) => {
                    this.buffer.extend_from_slice(&bytes);
                }
                Poll::Pending => {
                    return Poll::Pending;
                }
                Poll::Ready(None) => {
                    if this.buffer.is_empty() {
                        return Poll::Ready(None);
                    } else {
                        // If there's something left in the buffer, it means we received a malformed message.
                        let err_text = String::from_utf8_lossy(this.buffer).into_owned();
                        let err = GeminiResponseStreamError::InvalidResposeFormat(format!(
                            "Stream ended with incomplete data in the buffer: {}",
                            err_text
                        ));
                        // Clear the buffer to avoid triggering the error again.
                        this.buffer.clear();
                        return Poll::Ready(Some(Err(err)));
                    }
                }
                Poll::Ready(Some(Err(e))) => {
                    return Poll::Ready(Some(Err(GeminiResponseStreamError::ReqwestError(e))));
                }
            }
        }
    }
}
impl<F, T> ResponseStream<F, T>
where
    F: FnMut(&Session, GeminiResponse) -> T,
{
    pub(crate) fn new(
        response_stream: Box<
            dyn Stream<Item = Result<Bytes, reqwest::Error>> + Unpin + Send + 'static,
        >,
        session: Session,
        data_extractor: F,
    ) -> Self {
        Self {
            response_stream,
            session,
            data_extractor,
            buffer: Vec::new(),
        }
    }
    pub fn get_session(&self) -> &Session {
        &self.session
    }
    pub fn get_session_owned(self) -> Session {
        self.session
    }
}
pub type GeminiResponseStream =
    ResponseStream<fn(&Session, GeminiResponse) -> GeminiResponse, GeminiResponse>;
