use std::time::Duration;

use anyhow::{Context, Result, bail};
use serde_json::{Map, Value, json};
use tokio::time::sleep;
use url::Url;

use crate::{
    ai::{AgentStep, LlmProvider, Message, ToolCall},
    config::ResolvedProvider,
    term::format_api_error_summary,
    tools::ToolDef,
};

pub struct VertexProvider {
    client: reqwest::Client,
    config: ResolvedProvider,
}

impl VertexProvider {
    pub fn new(config: ResolvedProvider) -> Self {
        Self {
            client: reqwest::Client::new(),
            config,
        }
    }

    fn region(&self) -> String {
        if let Some(region) = self.config.extra_headers.get("x-appctl-vertex-region") {
            return region.clone();
        }

        Url::parse(&self.config.base_url)
            .ok()
            .and_then(|url| url.host_str().map(str::to_string))
            .and_then(|host| host.split('.').next().map(str::to_string))
            .filter(|part| !part.is_empty())
            .unwrap_or_else(|| "us-central1".to_string())
    }

    fn project_id(&self) -> Result<String> {
        self.config.auth_status.project_id.clone().context(
            "Vertex requires a Google Cloud project. Run `appctl auth provider login vertex`.",
        )
    }

    fn build_request(
        &self,
        access_token: &str,
        project_id: &str,
        region: &str,
        payload: &Value,
    ) -> reqwest::RequestBuilder {
        let mut request = self.client.post(format!(
            "{}/v1/projects/{}/locations/{}/publishers/google/models/{}:generateContent",
            self.config.base_url.trim_end_matches('/'),
            project_id,
            region,
            self.config.model
        ));
        request = request.bearer_auth(access_token);
        for (name, value) in &self.config.extra_headers {
            if name != "x-appctl-vertex-region" {
                request = request.header(name, value);
            }
        }
        request.json(payload)
    }
}

#[async_trait::async_trait]
impl LlmProvider for VertexProvider {
    async fn chat(&self, messages: &[Message], tools: &[ToolDef]) -> Result<AgentStep> {
        let access_token = self
            .config
            .auth
            .bearer_token()
            .context("Vertex requires Google ADC or another bearer token source")?;
        let region = self.region();
        let project_id = self.project_id()?;

        let payload = json!({
            "systemInstruction": system_instruction(messages),
            "contents": messages.iter().filter(|message| message.role != "system").map(serialize_message).collect::<Vec<_>>(),
            "tools": [
                {
                    "functionDeclarations": tools.iter().map(serialize_tool).collect::<Vec<_>>()
                }
            ]
        });

        let body = send_with_backoff(self, access_token, &project_id, &region, &payload).await?;
        let response: Value = serde_json::from_str(&body).with_context(|| {
            format!(
                "failed to parse Vertex AI response as JSON: {}",
                format_api_error_summary(&body)
            )
        })?;

        let candidate = response
            .pointer("/candidates/0/content/parts")
            .and_then(Value::as_array)
            .context("Vertex AI response missing candidates[0].content.parts")?;

        let mut tool_calls = Vec::new();
        let mut text = String::new();
        for part in candidate {
            if let Some(chunk) = part.get("text").and_then(Value::as_str) {
                text.push_str(chunk);
            }
            if let Some(call) = part.get("functionCall").and_then(Value::as_object) {
                tool_calls.push(ToolCall {
                    id: call
                        .get("id")
                        .and_then(Value::as_str)
                        .unwrap_or("tool")
                        .to_string(),
                    name: call
                        .get("name")
                        .and_then(Value::as_str)
                        .unwrap_or_default()
                        .to_string(),
                    arguments: call
                        .get("args")
                        .cloned()
                        .unwrap_or(Value::Object(Map::new())),
                });
            }
        }

        if !tool_calls.is_empty() {
            Ok(AgentStep::ToolCalls { calls: tool_calls })
        } else if text.is_empty() {
            Ok(AgentStep::Stop)
        } else {
            Ok(AgentStep::Message { content: text })
        }
    }
}

fn parse_retry_after_seconds(headers: &reqwest::header::HeaderMap) -> Option<u64> {
    headers
        .get(reqwest::header::RETRY_AFTER)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.trim().parse::<u64>().ok())
}

async fn send_with_backoff(
    provider: &VertexProvider,
    access_token: &str,
    project_id: &str,
    region: &str,
    payload: &Value,
) -> Result<String> {
    let max_retries = 3usize;
    let mut last_summary = String::new();

    for attempt in 0..=max_retries {
        let response = provider
            .build_request(access_token, project_id, region, payload)
            .send()
            .await
            .context("failed to call Vertex AI")?;
        let status = response.status();
        let headers = response.headers().clone();
        let body = response
            .text()
            .await
            .context("failed to read Vertex AI response body")?;

        if status.is_success() {
            return Ok(body);
        }

        let summary = format_api_error_summary(&body);
        if status.as_u16() != 429 {
            bail!("Vertex AI returned {}: {}", status, summary);
        }

        last_summary = summary;
        if attempt == max_retries {
            bail!(
                "Vertex AI returned 429 Too Many Requests after {} retries: {}. Try again shortly, switch to a lower-quota model like `gemini-1.5-flash`, or configure another provider for this app dir.",
                max_retries,
                last_summary
            );
        }

        let wait_secs = parse_retry_after_seconds(&headers).unwrap_or(1u64 << (attempt + 1));
        sleep(Duration::from_secs(wait_secs)).await;
    }

    bail!("Vertex AI rate-limited the request: {}", last_summary)
}

fn system_instruction(messages: &[Message]) -> Value {
    match messages.iter().find(|message| message.role == "system") {
        Some(message) => json!({
            "parts": [{ "text": message.content }]
        }),
        None => Value::Null,
    }
}

fn serialize_message(message: &Message) -> Value {
    match message.role.as_str() {
        "assistant" if !message.tool_calls.is_empty() => json!({
            "role": "model",
            "parts": message.tool_calls.iter().map(|call| json!({
                "functionCall": {
                    "name": call.name,
                    "args": call.arguments,
                }
            })).collect::<Vec<_>>()
        }),
        "assistant" => json!({
            "role": "model",
            "parts": [{ "text": message.content }]
        }),
        "tool" => {
            let response = serde_json::from_str::<Value>(&message.content)
                .unwrap_or_else(|_| Value::String(message.content.clone()));
            json!({
                "role": "user",
                "parts": [{
                    "functionResponse": {
                        "name": message.tool_name.clone().unwrap_or_default(),
                        "response": response,
                    }
                }]
            })
        }
        _ => json!({
            "role": "user",
            "parts": [{ "text": message.content }]
        }),
    }
}

fn serialize_tool(tool: &ToolDef) -> Value {
    json!({
        "name": tool.name,
        "description": tool.description,
        "parameters": tool.input_schema,
    })
}
