use crate::translator::{
    ModelId, ProviderCredentials, RewriteMode, RewriteProvider, RewriteRequest, RewriteResult,
    RewriteStream,
};
use async_trait::async_trait;
use futures_util::{StreamExt, future};
use regex::Regex;
use serde_json::json;

#[derive(Debug, Clone)]
struct ProtectedUrl {
    token: String,
    value: String,
}

#[derive(Debug, Clone)]
struct ProtectedText {
    text: String,
    urls: Vec<ProtectedUrl>,
}

pub struct GeminiProvider;

impl GeminiProvider {
    pub fn new() -> Result<Self, String> {
        Ok(Self)
    }

    pub(crate) fn build_prompt(text: &str, mode: RewriteMode) -> String {
        let mode_instruction = match mode {
            RewriteMode::Clean => {
                "Fix punctuation, capitalization, and obvious typos. Stay as close to the original sentence structure as possible."
            }
            RewriteMode::Readable => {
                "Improve the flow and sentence boundaries. Break up run-on sentences. Ensure the density of Information remains high, but make it easier for a third party to digest without losing the author's personality."
            }
            RewriteMode::Formal => {
                "Translate the core message into professional communication. Maintain accuracy and factual precision. Use clear, direct language suitable for a work context."
            }
            RewriteMode::Concise => {
                "Strip away filler words and repetitions. Keep the core meaning and tone intact but reduce the word count."
            }
            RewriteMode::PreserveVoice => {
                "Perform the absolute minimum interventions required for clarity. Fix only the most egregious errors. Preserve the raw pacing and stylistic quirks entirely."
            }
        };

        format!(
            "You are a silent translation layer. Your task is to transform raw, stream-of-consciousness writing into a polished, readable version while strictly preserving the author's unique voice, tone, and intent.\n\nRULES:\n- DO NOT repeat the original input text.\n- DO NOT add any preamble or postamble (for example: 'Here is your text').\n- DO NOT summarize.\n- DO NOT invent new facts.\n- DO NOT use generic AI-sounding phrases.\n- Preserve uncertainty, intensity, and factual precision.\n- Preserve every protected token exactly as written.\n- Preserve all links and URLs exactly. Never remove, shorten, rewrite, or relabel them.\n- OUTPUT ONLY the final translated text.\n\nTARGET MODE: {}\n\nRAW INPUT:\n{}\n\nTRANSLATED OUTPUT:",
            mode_instruction, text
        )
    }

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

    fn request_body(prompt: String) -> serde_json::Value {
        json!({
            "system_instruction": {
                "parts": [{
                    "text": "You are a fidelity-preserving communication translator. Improve readability while preserving the author's meaning, voice, uncertainty, emotional tone, and any protected tokens or links exactly."
                }]
            },
            "contents": [{
                "role": "user",
                "parts": [{"text": prompt}]
            }],
            "generationConfig": {
                "temperature": 0.1,
                "responseMimeType": "text/plain"
            }
        })
    }

    fn protect_urls(text: &str) -> ProtectedText {
        let url_regex = Regex::new(r"https?://\S+").expect("valid URL protection regex");
        let mut protected_text = text.to_string();
        let mut urls = Vec::new();
        let matches: Vec<_> = url_regex.find_iter(text).collect();

        for (index, capture) in matches.into_iter().enumerate().rev() {
            let token = format!("[[THISPER_URL_{index}]]");
            urls.push(ProtectedUrl {
                token: token.clone(),
                value: capture.as_str().to_string(),
            });
            protected_text.replace_range(capture.range(), &token);
        }

        urls.reverse();

        ProtectedText {
            text: protected_text,
            urls,
        }
    }

    fn restore_urls(text: &str, protected: &ProtectedText) -> RewriteResult<String> {
        let restored = protected
            .urls
            .iter()
            .fold(text.to_string(), |acc, protected_url| {
                acc.replace(&protected_url.token, &protected_url.value)
            });

        let missing_urls: Vec<&str> = protected
            .urls
            .iter()
            .filter_map(|protected_url| {
                if restored.contains(&protected_url.value) {
                    None
                } else {
                    Some(protected_url.value.as_str())
                }
            })
            .collect();

        if missing_urls.is_empty() {
            Ok(restored)
        } else {
            Err(format!(
                "Rewrite removed protected URLs. Missing: {}",
                missing_urls.join(", ")
            ))
        }
    }

    fn finalize_rewrite(text: &str, protected: &ProtectedText) -> RewriteResult<String> {
        Ok(Self::restore_urls(text.trim(), protected)?
            .trim()
            .to_string())
    }

    pub(crate) fn parse_sse_chunk(buffer: &mut String, chunk: &str) -> Vec<RewriteResult<String>> {
        buffer.push_str(&chunk.replace("\r\n", "\n"));

        let mut events = Vec::new();
        while let Some(event_end) = buffer.find("\n\n") {
            let raw_event = buffer[..event_end].to_string();
            buffer.drain(..event_end + 2);

            if raw_event.trim().is_empty() {
                continue;
            }

            events.push(Self::parse_sse_event(&raw_event).and_then(|maybe_text| {
                maybe_text
                    .filter(|text| !text.is_empty())
                    .ok_or_else(|| "__skip__".to_string())
            }));
        }

        events
            .into_iter()
            .filter_map(|result| match result {
                Err(err) if err == "__skip__" => None,
                other => Some(other),
            })
            .collect()
    }

    fn parse_sse_event(event: &str) -> RewriteResult<Option<String>> {
        let mut data_lines = Vec::new();
        for line in event.lines() {
            let trimmed = line.trim();
            if trimmed.is_empty() || trimmed.starts_with(':') || trimmed.starts_with("event:") {
                continue;
            }

            if let Some(data) = trimmed.strip_prefix("data:") {
                data_lines.push(data.trim_start());
            }
        }

        if data_lines.is_empty() {
            return Ok(None);
        }

        let payload = data_lines.join("\n");
        if payload == "[DONE]" {
            return Ok(None);
        }

        Self::extract_text_from_response(&payload).map(Some)
    }

    pub(crate) fn extract_text_from_response(payload: &str) -> RewriteResult<String> {
        let value: serde_json::Value = serde_json::from_str(payload)
            .map_err(|err| format!("Failed to parse Gemini stream event: {err}"))?;

        let mut text = String::new();
        if let Some(candidates) = value["candidates"].as_array() {
            for candidate in candidates {
                if let Some(parts) = candidate["content"]["parts"].as_array() {
                    for part in parts {
                        if let Some(part_text) = part["text"].as_str() {
                            text.push_str(part_text);
                        }
                    }
                }
            }
        }

        Ok(text)
    }
}

#[async_trait]
impl RewriteProvider for GeminiProvider {
    async fn rewrite(
        &self,
        request: &RewriteRequest,
        credentials: &ProviderCredentials,
    ) -> RewriteResult<String> {
        if credentials.api_key.is_empty() {
            return Err("Gemini API key is not configured. Please set it in Settings.".to_string());
        }

        let protected = Self::protect_urls(&request.text);
        let client = reqwest::Client::new();
        let prompt = Self::build_prompt(&protected.text, request.mode);
        let url = format!(
            "https://generativelanguage.googleapis.com/v1beta/{}:generateContent",
            Self::model_name(&request.model)
        );
        let body = Self::request_body(prompt);

        let res = client
            .post(&url)
            .header("x-goog-api-key", &credentials.api_key)
            .json(&body)
            .send()
            .await
            .map_err(|err| format!("Failed to reach Gemini API: {err}"))?;

        if !res.status().is_success() {
            let err_text = res
                .text()
                .await
                .unwrap_or_else(|_| "Unknown API error".to_string());
            return Err(format!("Gemini API error: {err_text}"));
        }

        let resp_json: serde_json::Value = res
            .json()
            .await
            .map_err(|err| format!("Failed to decode Gemini response: {err}"))?;

        let rewritten = Self::extract_text_from_response(&resp_json.to_string())?;
        Self::finalize_rewrite(&rewritten, &protected)
    }

    async fn rewrite_stream(
        &self,
        request: &RewriteRequest,
        credentials: &ProviderCredentials,
    ) -> RewriteResult<RewriteStream> {
        if credentials.api_key.is_empty() {
            return Err("Gemini API key is not configured. Please set it in Settings.".to_string());
        }

        let protected = Self::protect_urls(&request.text);
        let client = reqwest::Client::new();
        let prompt = Self::build_prompt(&protected.text, request.mode);
        let url = format!(
            "https://generativelanguage.googleapis.com/v1beta/{}:streamGenerateContent?alt=sse",
            Self::model_name(&request.model)
        );
        let body = Self::request_body(prompt);

        let res = client
            .post(&url)
            .header("x-goog-api-key", &credentials.api_key)
            .json(&body)
            .send()
            .await
            .map_err(|err| format!("Failed to reach Gemini API: {err}"))?;

        if !res.status().is_success() {
            let err_text = res
                .text()
                .await
                .unwrap_or_else(|_| "Unknown API error".to_string());
            return Err(format!("Gemini API error (stream): {err_text}"));
        }

        let collected = res
            .bytes_stream()
            .scan(String::new(), |buffer, chunk_res| {
                let items = match chunk_res {
                    Ok(chunk) => {
                        let chunk_text = String::from_utf8_lossy(&chunk);
                        Self::parse_sse_chunk(buffer, &chunk_text)
                    }
                    Err(err) => vec![Err(format!("Failed to read Gemini stream: {err}"))],
                };
                future::ready(Some(items))
            })
            .flat_map(futures_util::stream::iter)
            .collect::<Vec<_>>()
            .await;

        let mut rewritten = String::new();
        for item in collected {
            match item {
                Ok(part) => rewritten.push_str(&part),
                Err(err) => return Err(err),
            }
        }

        let restored = Self::finalize_rewrite(&rewritten, &protected)?;
        Ok(Box::pin(futures_util::stream::once(
            async move { Ok(restored) },
        )))
    }

    fn provider_name(&self) -> &'static str {
        "Google Gemini"
    }
}

#[cfg(test)]
mod tests {
    use super::GeminiProvider;
    use crate::translator::RewriteMode;

    #[test]
    fn prompt_includes_voice_preservation_rules() {
        let prompt = GeminiProvider::build_prompt("raw text", RewriteMode::PreserveVoice);
        assert!(prompt.contains("Preserve uncertainty, intensity, and factual precision."));
        assert!(prompt.contains("Preserve all links and URLs exactly."));
        assert!(prompt.contains(
            "TARGET MODE: Perform the absolute minimum interventions required for clarity."
        ));
    }

    #[test]
    fn url_protection_round_trips_links() {
        let protected = GeminiProvider::protect_urls(
            "See https://research.google/blog/turboquant-redefining-ai-efficiency-with-extreme-compression/ for details.",
        );

        assert!(protected.text.contains("[[THISPER_URL_0]]"));

        let restored =
            GeminiProvider::finalize_rewrite("See [[THISPER_URL_0]] for details.", &protected)
                .unwrap();

        assert!(restored.contains(
            "https://research.google/blog/turboquant-redefining-ai-efficiency-with-extreme-compression/"
        ));
    }

    #[test]
    fn finalize_rewrite_rejects_missing_protected_links() {
        let protected = GeminiProvider::protect_urls("Link: https://example.com and keep it.");

        let error = GeminiProvider::finalize_rewrite("Link removed.", &protected).unwrap_err();
        assert!(error.contains("Rewrite removed protected URLs."));
    }

    #[test]
    fn sse_parser_handles_chunk_boundaries() {
        let mut buffer = String::new();
        let first = "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hello";
        let second = "\"}]}}]}\n\ndata: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\" world\"}]}}]}\n\n";

        let first_results = GeminiProvider::parse_sse_chunk(&mut buffer, first);
        let second_results = GeminiProvider::parse_sse_chunk(&mut buffer, second);

        assert!(first_results.is_empty());
        assert_eq!(second_results.len(), 2);
        assert_eq!(second_results[0].as_ref().unwrap(), "Hello");
        assert_eq!(second_results[1].as_ref().unwrap(), " world");
    }

    #[test]
    fn extract_text_joins_all_text_parts() {
        let payload =
            r#"{"candidates":[{"content":{"parts":[{"text":"Hello"},{"text":" world"}]}}]}"#;
        assert_eq!(
            GeminiProvider::extract_text_from_response(payload).unwrap(),
            "Hello world"
        );
    }
}
