//! Google Vertex AI Gemini provider.
//!
//! Vertex exposes Gemini through `projects/{project}/locations/{location}/...`
//! routes and Google OAuth bearer tokens. The request mapper keeps Harn's
//! canonical OpenAI-like message list at the boundary and emits Vertex
//! `generateContent` JSON.

use crate::llm::api::{DeltaSender, LlmRequestPayload, LlmResult};
use crate::llm::provider::{LlmProvider, LlmProviderChat};
use crate::llm::providers::common::{
    apply_provider_overrides, google_function_declaration_tools, maybe_emit_delta,
    request_text_content, vm_err,
};
use crate::url_encoding::percent_encode_component;
use crate::value::VmError;

pub(crate) struct VertexProvider;

#[derive(Debug, serde::Deserialize)]
struct ServiceAccountKey {
    client_email: String,
    private_key: String,
    #[serde(default)]
    token_uri: Option<String>,
    #[serde(default)]
    project_id: Option<String>,
}

#[derive(Debug, serde::Serialize)]
struct ServiceAccountClaims<'a> {
    iss: &'a str,
    scope: &'a str,
    aud: &'a str,
    exp: usize,
    iat: usize,
}

impl VertexProvider {
    pub(crate) fn build_request_body(request: &LlmRequestPayload) -> serde_json::Value {
        let mut contents = Vec::new();
        let mut system_parts = Vec::new();
        if let Some(system) = request.system.as_deref() {
            if !system.is_empty() {
                system_parts.push(serde_json::json!({"text": system}));
            }
        }
        for message in &request.messages {
            let role = match message.get("role").and_then(|value| value.as_str()) {
                Some("assistant") => "model",
                Some("system") => {
                    let text = request_text_content(message);
                    if !text.is_empty() {
                        system_parts.push(serde_json::json!({"text": text}));
                    }
                    continue;
                }
                _ => "user",
            };
            let text = request_text_content(message);
            if text.is_empty() {
                continue;
            }
            contents.push(serde_json::json!({
                "role": role,
                "parts": [{"text": text}],
            }));
        }
        if let Some(prefill) = request.prefill.as_deref() {
            if !prefill.is_empty() {
                contents.push(serde_json::json!({
                    "role": "model",
                    "parts": [{"text": prefill}],
                }));
            }
        }

        let mut body = serde_json::json!({ "contents": contents });
        if !system_parts.is_empty() {
            body["systemInstruction"] = serde_json::json!({ "parts": system_parts });
        }
        let mut generation = serde_json::Map::new();
        if request.max_tokens > 0 {
            generation.insert(
                "maxOutputTokens".to_string(),
                serde_json::json!(request.max_tokens),
            );
        }
        if let Some(temp) = request.temperature {
            generation.insert("temperature".to_string(), serde_json::json!(temp));
        }
        if let Some(top_p) = request.top_p {
            generation.insert("topP".to_string(), serde_json::json!(top_p));
        }
        if let Some(stop) = request.stop.as_ref() {
            generation.insert("stopSequences".to_string(), serde_json::json!(stop));
        }
        // The modern `output_format` is the source of truth; Vertex previously
        // read only the legacy `response_format`/`json_schema` mirror, which is
        // never backfilled from a nested `output_format.schema`, so a call using
        // `output_format: {kind: "json_schema", schema}` silently produced no
        // structured-output directive. Honor `output_format` first, falling back
        // to the legacy mirror for callers that only set those.
        match &request.output_format {
            crate::llm::api::OutputFormat::JsonObject => {
                generation.insert(
                    "responseMimeType".to_string(),
                    serde_json::json!("application/json"),
                );
            }
            crate::llm::api::OutputFormat::JsonSchema { schema, .. } => {
                generation.insert(
                    "responseMimeType".to_string(),
                    serde_json::json!("application/json"),
                );
                generation.insert("responseSchema".to_string(), schema.clone());
            }
            crate::llm::api::OutputFormat::Text => {
                if request.response_format.as_deref() == Some("json") {
                    generation.insert(
                        "responseMimeType".to_string(),
                        serde_json::json!("application/json"),
                    );
                    if let Some(schema) = request.json_schema.as_ref() {
                        generation.insert("responseSchema".to_string(), schema.clone());
                    }
                }
            }
        }
        if !generation.is_empty() {
            body["generationConfig"] = serde_json::Value::Object(generation);
        }
        if let Some(tools) = google_function_declaration_tools(request.native_tools.as_deref()) {
            body["tools"] = tools;
        }
        body
    }

    pub(crate) fn endpoint_url(request: &LlmRequestPayload) -> Result<String, VmError> {
        let project = std::env::var("VERTEX_AI_PROJECT")
            .or_else(|_| std::env::var("GOOGLE_CLOUD_PROJECT"))
            .or_else(|_| service_account_project())
            .map_err(|_| vm_err("Vertex AI project is not configured; set VERTEX_AI_PROJECT"))?;
        let location = std::env::var("VERTEX_AI_LOCATION")
            .ok()
            .filter(|value| !value.trim().is_empty())
            .unwrap_or_else(|| "us-central1".to_string());
        let pdef = crate::llm_config::provider_config("vertex");
        let base_url = pdef
            .as_ref()
            .map(crate::llm_config::resolve_base_url)
            .unwrap_or_else(|| "https://aiplatform.googleapis.com/v1".to_string());
        let base_url = base_url.trim_end_matches('/');
        let wire_model = crate::llm_config::wire_model_id(&request.model);
        let model = if wire_model.starts_with("projects/") {
            wire_model
        } else {
            format!(
                "projects/{}/locations/{}/publishers/google/models/{}",
                percent_encode_component(&project),
                percent_encode_component(&location),
                percent_encode_component(&wire_model)
            )
        };
        Ok(format!("{base_url}/{model}:generateContent"))
    }

    async fn bearer_token(api_key: &str) -> Result<String, VmError> {
        for env_name in ["VERTEX_AI_ACCESS_TOKEN", "GOOGLE_OAUTH_ACCESS_TOKEN"] {
            if let Ok(token) = std::env::var(env_name) {
                if !token.trim().is_empty() {
                    return Ok(token);
                }
            }
        }
        if !api_key.trim().is_empty() && !api_key.trim().starts_with('/') {
            return Ok(api_key.to_string());
        }
        let key_path = std::env::var("GOOGLE_APPLICATION_CREDENTIALS")
            .ok()
            .filter(|value| !value.trim().is_empty())
            .or_else(|| (!api_key.trim().is_empty()).then(|| api_key.to_string()))
            .ok_or_else(|| {
                vm_err(
                    "Missing Vertex AI credentials: set VERTEX_AI_ACCESS_TOKEN or GOOGLE_APPLICATION_CREDENTIALS",
                )
            })?;
        exchange_service_account_token(&key_path).await
    }

    pub(crate) async fn chat_impl(
        &self,
        request: &LlmRequestPayload,
        delta_tx: Option<DeltaSender>,
    ) -> Result<LlmResult, VmError> {
        let url = Self::endpoint_url(request)?;
        let token = Self::bearer_token(&request.api_key).await?;
        let mut body = Self::build_request_body(request);
        apply_provider_overrides(&mut body, request.provider_overrides.as_ref());
        let response = crate::llm::shared_blocking_client()
            .post(url)
            .header("Content-Type", "application/json")
            .header("Authorization", format!("Bearer {token}"))
            .timeout(std::time::Duration::from_secs(request.resolve_timeout()))
            .json(&body)
            .send()
            .await
            .map_err(|error| vm_err(format!("vertex API error: {error}")))?;
        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(vm_err(
                crate::llm::api::classify_provider_http_error("vertex", status, None, &body)
                    .message,
            ));
        }
        let json: serde_json::Value = response
            .json()
            .await
            .map_err(|error| vm_err(format!("vertex response parse error: {error}")))?;
        let result = parse_vertex_response(&json, &request.model)?;
        maybe_emit_delta(delta_tx, &result.text);
        Ok(result)
    }
}

impl LlmProvider for VertexProvider {
    fn name(&self) -> &'static str {
        "vertex"
    }
}

impl LlmProviderChat for VertexProvider {
    fn chat<'a>(
        &'a self,
        request: &'a LlmRequestPayload,
        delta_tx: Option<DeltaSender>,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<LlmResult, VmError>> + 'a>> {
        Box::pin(self.chat_impl(request, delta_tx))
    }
}

fn parse_vertex_response(json: &serde_json::Value, model: &str) -> Result<LlmResult, VmError> {
    if let Some(error) = json["error"]["message"].as_str() {
        return Err(vm_err(format!("vertex API error: {error}")));
    }
    let mut result = crate::llm::providers::common::empty_result("vertex", model);
    if let Some(parts) = json["candidates"][0]["content"]["parts"].as_array() {
        for part in parts {
            if let Some(text) = part.get("text").and_then(|value| value.as_str()) {
                result.text.push_str(text);
                result.blocks.push(serde_json::json!({
                    "type": "output_text",
                    "text": text,
                    "visibility": "public",
                }));
            }
            if let Some(call) = part.get("functionCall") {
                let name = call
                    .get("name")
                    .and_then(|value| value.as_str())
                    .unwrap_or("")
                    .to_string();
                let args = call
                    .get("args")
                    .cloned()
                    .unwrap_or_else(|| serde_json::json!({}));
                let id = format!("vertex_tool_{}", result.tool_calls.len());
                result.tool_calls.push(serde_json::json!({
                    "id": id,
                    "name": name,
                    "arguments": args,
                }));
                result.blocks.push(serde_json::json!({
                    "type": "tool_call",
                    "id": id,
                    "name": name,
                    "arguments": args,
                    "visibility": "internal",
                }));
            }
        }
    }
    result.input_tokens = json["usageMetadata"]["promptTokenCount"]
        .as_i64()
        .unwrap_or(0);
    result.output_tokens = json["usageMetadata"]["candidatesTokenCount"]
        .as_i64()
        .unwrap_or(0)
        + json["usageMetadata"]["thoughtsTokenCount"]
            .as_i64()
            .unwrap_or(0);
    result.cache_read_tokens = json["usageMetadata"]["cachedContentTokenCount"]
        .as_i64()
        .unwrap_or(0);
    result.stop_reason = json["candidates"][0]["finishReason"]
        .as_str()
        .map(str::to_string);
    Ok(result)
}

fn service_account_project() -> Result<String, VmError> {
    let path = std::env::var("GOOGLE_APPLICATION_CREDENTIALS")
        .map_err(|_| vm_err("GOOGLE_APPLICATION_CREDENTIALS is not set"))?;
    let key = read_service_account_key(&path)?;
    key.project_id
        .filter(|project| !project.trim().is_empty())
        .ok_or_else(|| vm_err("service account JSON does not contain project_id"))
}

fn read_service_account_key(path: &str) -> Result<ServiceAccountKey, VmError> {
    let text = std::fs::read_to_string(path)
        .map_err(|error| vm_err(format!("failed to read service account JSON: {error}")))?;
    serde_json::from_str(&text)
        .map_err(|error| vm_err(format!("failed to parse service account JSON: {error}")))
}

async fn exchange_service_account_token(path: &str) -> Result<String, VmError> {
    let key = read_service_account_key(path)?;
    let token_uri = key
        .token_uri
        .as_deref()
        .unwrap_or("https://oauth2.googleapis.com/token");
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map_err(|error| vm_err(format!("system clock error: {error}")))?
        .as_secs() as usize;
    let claims = ServiceAccountClaims {
        iss: &key.client_email,
        scope: "https://www.googleapis.com/auth/cloud-platform",
        aud: token_uri,
        iat: now,
        exp: now + 3600,
    };
    let jwt = jsonwebtoken::encode(
        &jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256),
        &claims,
        &jsonwebtoken::EncodingKey::from_rsa_pem(key.private_key.as_bytes())
            .map_err(|error| vm_err(format!("invalid service account private key: {error}")))?,
    )
    .map_err(|error| vm_err(format!("failed to sign service account JWT: {error}")))?;
    let response = crate::llm::shared_utility_client()
        .post(token_uri)
        .form(&[
            ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"),
            ("assertion", jwt.as_str()),
        ])
        .send()
        .await
        .map_err(|error| vm_err(format!("service account token exchange failed: {error}")))?;
    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        return Err(vm_err(format!(
            "service account token exchange HTTP {status}: {body}"
        )));
    }
    let json: serde_json::Value = response.json().await.map_err(|error| {
        vm_err(format!(
            "service account token response parse error: {error}"
        ))
    })?;
    json["access_token"]
        .as_str()
        .map(str::to_string)
        .ok_or_else(|| vm_err("service account token response did not include access_token"))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::llm::api::{LlmRequestPayload, ThinkingConfig};
    use serde_json::json;

    #[test]
    fn build_request_maps_messages_to_generate_content() {
        let body = VertexProvider::build_request_body(&base_request());
        assert_eq!(body["systemInstruction"]["parts"][0]["text"], "be brief");
        assert_eq!(body["contents"][0]["role"], "user");
        assert_eq!(body["contents"][0]["parts"][0]["text"], "hello");
        assert_eq!(body["generationConfig"]["maxOutputTokens"], 32);
    }

    #[test]
    fn build_request_honors_modern_output_format_json_schema() {
        // Regression: Vertex read only the legacy response_format/json_schema
        // mirror, so a modern `output_format` was silently dropped.
        let mut request = base_request();
        request.output_format = crate::llm::api::OutputFormat::JsonSchema {
            schema: json!({"type": "object", "properties": {"answer": {"type": "string"}}}),
            strict: true,
        };
        let body = VertexProvider::build_request_body(&request);
        assert_eq!(
            body["generationConfig"]["responseMimeType"],
            "application/json"
        );
        assert_eq!(
            body["generationConfig"]["responseSchema"]["properties"]["answer"]["type"],
            "string"
        );
    }

    #[test]
    fn parse_response_extracts_text_usage_and_function_calls() {
        let response = json!({
            "candidates": [{
                "content": {"parts": [
                    {"text": "hi"},
                    {"functionCall": {"name": "lookup", "args": {"q": "x"}}}
                ]},
                "finishReason": "STOP"
            }],
            "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 4}
        });
        let result = parse_vertex_response(&response, "gemini-1.5-pro-002").expect("result");
        assert_eq!(result.text, "hi");
        assert_eq!(result.input_tokens, 3);
        assert_eq!(result.output_tokens, 4);
        assert_eq!(result.tool_calls[0]["name"], "lookup");
    }

    #[test]
    fn parse_response_folds_thinking_and_cache_tokens() {
        let response = json!({
            "candidates": [{
                "content": {"parts": [{"text": "done"}]},
                "finishReason": "STOP"
            }],
            "usageMetadata": {
                "promptTokenCount": 12,
                "candidatesTokenCount": 4,
                "thoughtsTokenCount": 9,
                "cachedContentTokenCount": 6
            }
        });
        let result = parse_vertex_response(&response, "gemini-2.5-pro").expect("result");
        assert_eq!(result.input_tokens, 12);
        // candidates(4) + thoughts(9) so thinking tokens are billed as output.
        assert_eq!(result.output_tokens, 13);
        assert_eq!(result.cache_read_tokens, 6);
    }

    fn base_request() -> LlmRequestPayload {
        LlmRequestPayload {
            provider: "vertex".to_string(),
            model: "gemini-1.5-pro-002".to_string(),
            region: None,
            api_key: String::new(),
            api_mode: crate::llm::api::LlmApiMode::ChatCompletions,
            fallback_chain: Vec::new(),
            route_fallbacks: Vec::new(),
            messages: vec![json!({"role": "user", "content": "hello"})],
            system: Some("be brief".to_string()),
            max_tokens: 32,
            temperature: Some(0.2),
            top_p: None,
            top_k: None,
            logprobs: false,
            top_logprobs: None,
            stop: None,
            seed: None,
            frequency_penalty: None,
            presence_penalty: None,
            fast: false,
            output_format: crate::llm::api::OutputFormat::Text,
            response_format: None,
            json_schema: None,
            output_schema: None,
            schema_stream_abort: false,
            thinking: ThinkingConfig::Disabled,
            anthropic_beta_features: Vec::new(),
            vision: false,
            native_tools: Some(vec![json!({
                "type": "function",
                "function": {
                    "name": "lookup",
                    "description": "Lookup",
                    "parameters": {"type": "object"}
                }
            })]),
            provider_tools: Vec::new(),
            tool_choice: None,
            cache: false,
            timeout: None,
            stream: false,
            provider_overrides: None,
            previous_response_id: None,
            store: None,
            background: None,
            truncation: None,
            compact: None,
            include: None,
            max_tool_calls: None,
            prefill: None,
            session_id: None,
            reminder_lifecycle: Vec::new(),
        }
    }
}
