//! Gemini runtime — drives Google Gemini API sessions as agent backends.
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};

use reqwest::Client;
use serde_json::json;
use tokio::sync::broadcast;

use crate::agent_log::AgentLogWriter;
use crate::http::mcp::dispatch::dispatch_tool_call;
use crate::slog;

use super::super::{AgentEvent, TokenUsage};
use super::{AgentRuntime, RuntimeContext, RuntimeResult, RuntimeStatus};

mod api;
mod mcp;

use api::{build_generate_content_request, build_system_instruction, parse_usage_metadata};
use mcp::convert_mcp_tools_to_gemini;

// ── Internal types ───────────────────────────────────────────────────

struct GeminiFunctionCall {
    name: String,
    args: serde_json::Value,
}

// ── Public runtime struct ────────────────────────────────────────────

/// Agent runtime that drives a Gemini model through the Google AI
/// `generateContent` REST API.
///
/// The runtime:
/// 1. Fetches MCP tool definitions from huskies' MCP server.
/// 2. Converts them to Gemini function-calling format.
/// 3. Sends the agent prompt + tools to the Gemini API.
/// 4. Executes any requested function calls via MCP `tools/call`.
/// 5. Loops until the model produces a text-only response or an error.
/// 6. Tracks token usage from the API response metadata.
pub struct GeminiRuntime {
    /// Whether a stop has been requested.
    cancelled: Arc<AtomicBool>,
}

impl GeminiRuntime {
    /// Create a new Gemini runtime instance.
    pub fn new() -> Self {
        Self {
            cancelled: Arc::new(AtomicBool::new(false)),
        }
    }
}

impl AgentRuntime for GeminiRuntime {
    async fn start(
        &self,
        ctx: RuntimeContext,
        tx: broadcast::Sender<AgentEvent>,
        event_log: Arc<Mutex<Vec<AgentEvent>>>,
        log_writer: Option<Arc<Mutex<AgentLogWriter>>>,
    ) -> Result<RuntimeResult, String> {
        let api_key = std::env::var("GOOGLE_AI_API_KEY").map_err(|_| {
            "GOOGLE_AI_API_KEY environment variable is not set. \
             Set it to your Google AI API key to use the Gemini runtime."
                .to_string()
        })?;

        let model = if ctx.command.starts_with("gemini") {
            // The pool puts the model into `command` for non-CLI runtimes,
            // but also check args for a --model flag.
            ctx.command.clone()
        } else {
            // Fall back to args: look for --model <value>
            ctx.args
                .iter()
                .position(|a| a == "--model")
                .and_then(|i| ctx.args.get(i + 1))
                .cloned()
                .unwrap_or_else(|| "gemini-2.5-pro".to_string())
        };

        let app_ctx = ctx
            .app_ctx
            .clone()
            .ok_or_else(|| "Gemini runtime requires app_ctx to be set".to_string())?;

        let client = Client::new();
        let cancelled = Arc::clone(&self.cancelled);

        // Step 1: Load MCP tool definitions and convert to Gemini format.
        let gemini_tools = convert_mcp_tools_to_gemini();

        // Step 2: Build the initial conversation contents.
        let system_instruction = build_system_instruction(&ctx);
        let mut contents: Vec<serde_json::Value> = vec![json!({
            "role": "user",
            "parts": [{ "text": ctx.prompt }]
        })];

        let mut total_usage = TokenUsage {
            input_tokens: 0,
            output_tokens: 0,
            cache_creation_input_tokens: 0,
            cache_read_input_tokens: 0,
            total_cost_usd: 0.0,
        };

        let emit = |event: AgentEvent| {
            super::super::pty::emit_event(
                event,
                &tx,
                &event_log,
                log_writer.as_ref().map(|w| w.as_ref()),
            );
        };

        emit(AgentEvent::Status {
            story_id: ctx.story_id.clone(),
            agent_name: ctx.agent_name.clone(),
            status: "running".to_string(),
        });

        // Step 3: Conversation loop.
        let mut turn = 0u32;
        let max_turns = 200; // Safety limit

        loop {
            if cancelled.load(Ordering::Relaxed) {
                emit(AgentEvent::Error {
                    story_id: ctx.story_id.clone(),
                    agent_name: ctx.agent_name.clone(),
                    message: "Agent was stopped by user".to_string(),
                });
                return Ok(RuntimeResult {
                    session_id: None,
                    token_usage: Some(total_usage),
                    exit_ok: true,
                    aborted_signal: false,
                    rate_limit_exit: false,
                    rate_limit_reset_at: None,
                });
            }

            turn += 1;
            if turn > max_turns {
                emit(AgentEvent::Error {
                    story_id: ctx.story_id.clone(),
                    agent_name: ctx.agent_name.clone(),
                    message: format!("Exceeded maximum turns ({max_turns})"),
                });
                return Ok(RuntimeResult {
                    session_id: None,
                    token_usage: Some(total_usage),
                    exit_ok: true,
                    aborted_signal: false,
                    rate_limit_exit: false,
                    rate_limit_reset_at: None,
                });
            }

            slog!(
                "[gemini] Turn {turn} for {}:{}",
                ctx.story_id,
                ctx.agent_name
            );

            let request_body =
                build_generate_content_request(&system_instruction, &contents, &gemini_tools);

            let url = format!(
                "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
            );

            let response = client
                .post(&url)
                .json(&request_body)
                .send()
                .await
                .map_err(|e| format!("Gemini API request failed: {e}"))?;

            let status = response.status();
            let body: serde_json::Value = response
                .json()
                .await
                .map_err(|e| format!("Failed to parse Gemini API response: {e}"))?;

            if !status.is_success() {
                let error_msg = body["error"]["message"]
                    .as_str()
                    .unwrap_or("Unknown API error");
                let err = format!("Gemini API error ({status}): {error_msg}");
                emit(AgentEvent::Error {
                    story_id: ctx.story_id.clone(),
                    agent_name: ctx.agent_name.clone(),
                    message: err.clone(),
                });
                return Err(err);
            }

            // Accumulate token usage.
            if let Some(usage) = parse_usage_metadata(&body) {
                total_usage.input_tokens += usage.input_tokens;
                total_usage.output_tokens += usage.output_tokens;
            }

            // Extract the candidate response.
            let candidate = body["candidates"]
                .as_array()
                .and_then(|c| c.first())
                .ok_or_else(|| "No candidates in Gemini response".to_string())?;

            let parts = candidate["content"]["parts"]
                .as_array()
                .ok_or_else(|| "No parts in Gemini response candidate".to_string())?;

            // Check finish reason.
            let finish_reason = candidate["finishReason"].as_str().unwrap_or("");

            // Separate text parts and function call parts.
            let mut text_parts: Vec<String> = Vec::new();
            let mut function_calls: Vec<GeminiFunctionCall> = Vec::new();

            for part in parts {
                if let Some(text) = part["text"].as_str() {
                    text_parts.push(text.to_string());
                }
                if let Some(fc) = part.get("functionCall")
                    && let (Some(name), Some(args)) = (fc["name"].as_str(), fc.get("args"))
                {
                    function_calls.push(GeminiFunctionCall {
                        name: name.to_string(),
                        args: args.clone(),
                    });
                }
            }

            // Emit any text output.
            for text in &text_parts {
                if !text.is_empty() {
                    emit(AgentEvent::Output {
                        story_id: ctx.story_id.clone(),
                        agent_name: ctx.agent_name.clone(),
                        text: text.clone(),
                    });
                }
            }

            // If no function calls, the model is done.
            if function_calls.is_empty() {
                emit(AgentEvent::Done {
                    story_id: ctx.story_id.clone(),
                    agent_name: ctx.agent_name.clone(),
                    session_id: None,
                });
                return Ok(RuntimeResult {
                    session_id: None,
                    token_usage: Some(total_usage),
                    exit_ok: true,
                    aborted_signal: false,
                    rate_limit_exit: false,
                    rate_limit_reset_at: None,
                });
            }

            // Add the model's response to the conversation.
            let model_parts: Vec<serde_json::Value> = parts.to_vec();
            contents.push(json!({
                "role": "model",
                "parts": model_parts
            }));

            // Execute function calls via MCP and build response parts.
            let mut response_parts: Vec<serde_json::Value> = Vec::new();

            for fc in &function_calls {
                if cancelled.load(Ordering::Relaxed) {
                    break;
                }

                slog!(
                    "[gemini] Calling MCP tool '{}' for {}:{}",
                    fc.name,
                    ctx.story_id,
                    ctx.agent_name
                );

                emit(AgentEvent::Output {
                    story_id: ctx.story_id.clone(),
                    agent_name: ctx.agent_name.clone(),
                    text: format!("\n[Tool call: {}]\n", fc.name),
                });

                let tool_result = dispatch_tool_call(&fc.name, fc.args.clone(), &app_ctx).await;

                let response_value = match &tool_result {
                    Ok(result) => {
                        emit(AgentEvent::Output {
                            story_id: ctx.story_id.clone(),
                            agent_name: ctx.agent_name.clone(),
                            text: format!("[Tool result: {} chars]\n", result.len()),
                        });
                        json!({ "result": result })
                    }
                    Err(e) => {
                        emit(AgentEvent::Output {
                            story_id: ctx.story_id.clone(),
                            agent_name: ctx.agent_name.clone(),
                            text: format!("[Tool error: {e}]\n"),
                        });
                        json!({ "error": e })
                    }
                };

                response_parts.push(json!({
                    "functionResponse": {
                        "name": fc.name,
                        "response": response_value
                    }
                }));
            }

            // Add function responses to the conversation.
            contents.push(json!({
                "role": "user",
                "parts": response_parts
            }));

            // If the model indicated it's done despite having function calls,
            // respect the finish reason.
            if finish_reason == "STOP" && function_calls.is_empty() {
                break;
            }
        }

        emit(AgentEvent::Done {
            story_id: ctx.story_id.clone(),
            agent_name: ctx.agent_name.clone(),
            session_id: None,
        });

        Ok(RuntimeResult {
            session_id: None,
            token_usage: Some(total_usage),
            exit_ok: true,
            aborted_signal: false,
            rate_limit_exit: false,
            rate_limit_reset_at: None,
        })
    }

    fn stop(&self) {
        self.cancelled.store(true, Ordering::Relaxed);
    }

    fn get_status(&self) -> RuntimeStatus {
        if self.cancelled.load(Ordering::Relaxed) {
            RuntimeStatus::Failed
        } else {
            RuntimeStatus::Idle
        }
    }
}

// ── Tests ────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::http::context::AppContext;

    fn test_app_ctx() -> Arc<AppContext> {
        let tmp = tempfile::tempdir().unwrap();
        Arc::new(AppContext::new_test(tmp.path().to_path_buf()))
    }

    #[test]
    fn gemini_runtime_stop_sets_cancelled() {
        let runtime = GeminiRuntime::new();
        assert_eq!(runtime.get_status(), RuntimeStatus::Idle);
        runtime.stop();
        assert_eq!(runtime.get_status(), RuntimeStatus::Failed);
    }

    #[test]
    fn model_extraction_from_command() {
        // When command starts with "gemini", use it as model name
        let ctx = RuntimeContext {
            story_id: "1".to_string(),
            agent_name: "coder".to_string(),
            command: "gemini-2.5-pro".to_string(),
            args: vec![],
            prompt: "test".to_string(),
            cwd: "/tmp".to_string(),
            inactivity_timeout_secs: 300,
            app_ctx: Some(test_app_ctx()),
            session_id_to_resume: None,
            fresh_prompt: None,
            project_root: std::path::PathBuf::from("/tmp/project"),
            model: None,
        };

        // The model extraction logic is inside start(), but we test the
        // condition here.
        assert!(ctx.command.starts_with("gemini"));
    }
}
