use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::fmt::{Display, Formatter};
use std::pin::Pin;
use std::str::FromStr;

use futures_util::Stream;

pub type RewriteResult<T> = Result<T, String>;
pub type RewriteStream = Pin<Box<dyn Stream<Item = RewriteResult<String>> + Send>>;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RewriteMode {
    Clean,
    Readable,
    Formal,
    Concise,
    PreserveVoice,
}

impl FromStr for RewriteMode {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value {
            "Clean" => Ok(Self::Clean),
            "Readable" => Ok(Self::Readable),
            "Formal" => Ok(Self::Formal),
            "Concise" => Ok(Self::Concise),
            "Preserve Voice" => Ok(Self::PreserveVoice),
            _ => Err(format!("Unknown mode: {value}")),
        }
    }
}

impl Display for RewriteMode {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let value = match self {
            Self::Clean => "Clean",
            Self::Readable => "Readable",
            Self::Formal => "Formal",
            Self::Concise => "Concise",
            Self::PreserveVoice => "Preserve Voice",
        };

        f.write_str(value)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum GeminiModel {
    Gemini20FlashLite,
    Gemini31FlashLitePreview,
    GeminiFlashLiteLatest,
    Gemini25Flash,
}

impl GeminiModel {
    pub fn api_name(self) -> &'static str {
        match self {
            Self::Gemini20FlashLite => "models/gemini-2.0-flash-lite",
            Self::Gemini31FlashLitePreview => "models/gemini-3.1-flash-lite-preview",
            Self::GeminiFlashLiteLatest => "models/gemini-flash-lite-latest",
            Self::Gemini25Flash => "models/gemini-2.5-flash",
        }
    }
}

impl Default for GeminiModel {
    fn default() -> Self {
        Self::Gemini20FlashLite
    }
}

impl FromStr for GeminiModel {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value {
            "models/gemini-2.0-flash-lite" => Ok(Self::Gemini20FlashLite),
            "models/gemini-3.1-flash-lite-preview" => Ok(Self::Gemini31FlashLitePreview),
            "models/gemini-flash-lite-latest" => Ok(Self::GeminiFlashLiteLatest),
            "models/gemini-2.5-flash" => Ok(Self::Gemini25Flash),
            _ => Err(format!("Unknown model: {value}")),
        }
    }
}

impl Display for GeminiModel {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.api_name())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum ModelId {
    Gemini(GeminiModel),
}

impl Default for ModelId {
    fn default() -> Self {
        Self::Gemini(GeminiModel::default())
    }
}

impl ModelId {
    pub fn api_name(&self) -> &'static str {
        match self {
            Self::Gemini(model) => model.api_name(),
        }
    }
}

impl Display for ModelId {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.api_name())
    }
}

impl FromStr for ModelId {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        Ok(Self::Gemini(GeminiModel::from_str(value)?))
    }
}

#[derive(Debug, Clone)]
pub struct RewriteRequest {
    pub text: String,
    pub mode: RewriteMode,
    pub model: ModelId,
}

#[derive(Debug, Clone)]
pub struct ProviderCredentials {
    pub api_key: String,
}

#[async_trait]
pub trait RewriteProvider: Send + Sync {
    async fn rewrite(
        &self,
        request: &RewriteRequest,
        credentials: &ProviderCredentials,
    ) -> RewriteResult<String>;

    async fn rewrite_stream(
        &self,
        request: &RewriteRequest,
        credentials: &ProviderCredentials,
    ) -> RewriteResult<RewriteStream>;

    fn provider_name(&self) -> &'static str;
}

#[cfg(test)]
mod tests {
    use super::{GeminiModel, ModelId, RewriteMode};
    use std::str::FromStr;

    #[test]
    fn rewrite_mode_parses_known_values() {
        assert_eq!(RewriteMode::from_str("Clean").unwrap(), RewriteMode::Clean);
        assert_eq!(
            RewriteMode::from_str("Preserve Voice").unwrap(),
            RewriteMode::PreserveVoice
        );
    }

    #[test]
    fn rewrite_mode_rejects_unknown_values() {
        assert!(RewriteMode::from_str("Aggressive").is_err());
    }

    #[test]
    fn model_id_parses_known_values() {
        assert_eq!(
            ModelId::from_str("models/gemini-2.5-flash").unwrap(),
            ModelId::Gemini(GeminiModel::Gemini25Flash)
        );
    }
}
