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

use reqwest::Client;
use serde_json::{Value, json};
use tokio::sync::broadcast;

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

use super::super::{AgentEvent, TokenUsage};
use super::api_common::{
    CancellationFlag, build_system_text, check_loop_guard, clean_schema_properties, done_result,
    extract_model, start_conversation_loop,
};
use super::{AgentRuntime, RuntimeContext, RuntimeResult, RuntimeStatus};

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

/// Agent runtime that drives an OpenAI model (GPT-4o, o3, etc.) through
/// the OpenAI Chat Completions API.
///
/// The runtime:
/// 1. Fetches MCP tool definitions from huskies' MCP server.
/// 2. Converts them to OpenAI function-calling format.
/// 3. Sends the agent prompt + tools to the Chat Completions API.
/// 4. Executes any requested tool calls via MCP `tools/call`.
/// 5. Loops until the model produces a response with no tool calls.
/// 6. Tracks token usage from the API response.
pub struct OpenAiRuntime {
    /// Whether a stop has been requested.
    cancelled: CancellationFlag,
}

impl OpenAiRuntime {
    /// Create a new OpenAI runtime instance.
    pub fn new() -> Self {
        Self {
            cancelled: CancellationFlag::new(),
        }
    }
}

impl AgentRuntime for OpenAiRuntime {
    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("OPENAI_API_KEY").map_err(|_| {
            "OPENAI_API_KEY environment variable is not set. \
             Set it to your OpenAI API key to use the OpenAI runtime."
                .to_string()
        })?;

        let model = extract_model(
            &ctx,
            |c| c.starts_with("gpt") || c.starts_with("o"),
            "gpt-4o",
        );

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

        let client = Client::new();
        let cancelled = self.cancelled.handle();

        // Step 1: Fetch MCP tool definitions and convert to OpenAI format.
        let openai_tools = convert_mcp_tools_to_openai();

        // Step 2: Build the initial conversation messages.
        let system_text = build_system_text(&ctx);
        let mut messages: Vec<Value> = vec![
            json!({ "role": "system", "content": system_text }),
            json!({ "role": "user", "content": ctx.prompt }),
        ];

        let (emit, mut total_usage, mut turn) =
            start_conversation_loop(&ctx, tx, event_log, log_writer);

        // Step 3: Conversation loop.
        loop {
            if let Some(result) = check_loop_guard(&ctx, &cancelled, &mut turn, &total_usage, &emit)
            {
                return Ok(result);
            }

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

            let mut request_body = json!({
                "model": model,
                "messages": messages,
                "temperature": 0.2,
            });

            if !openai_tools.is_empty() {
                request_body["tools"] = json!(openai_tools);
            }

            let response = client
                .post("https://api.openai.com/v1/chat/completions")
                .bearer_auth(&api_key)
                .json(&request_body)
                .send()
                .await
                .map_err(|e| format!("OpenAI API request failed: {e}"))?;

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

            if !status.is_success() {
                let error_msg = body["error"]["message"]
                    .as_str()
                    .unwrap_or("Unknown API error");
                let err = format!("OpenAI 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(&body) {
                total_usage.input_tokens += usage.input_tokens;
                total_usage.output_tokens += usage.output_tokens;
            }

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

            let message = &choice["message"];
            let content = message["content"].as_str().unwrap_or("");

            // Emit any text content.
            if !content.is_empty() {
                emit(AgentEvent::Output {
                    story_id: ctx.story_id.clone(),
                    agent_name: ctx.agent_name.clone(),
                    text: content.to_string(),
                });
            }

            // Check for tool calls.
            let tool_calls = message["tool_calls"].as_array();

            if tool_calls.is_none() || tool_calls.is_some_and(|tc| tc.is_empty()) {
                // No tool calls — model is done.
                return Ok(done_result(&ctx, &emit, total_usage));
            }

            let tool_calls = tool_calls.unwrap();

            // Add the assistant message (with tool_calls) to the conversation.
            messages.push(message.clone());

            // Execute each tool call via MCP and add results.
            for tc in tool_calls {
                if cancelled.load(Ordering::Relaxed) {
                    break;
                }

                let call_id = tc["id"].as_str().unwrap_or("");
                let function = &tc["function"];
                let tool_name = function["name"].as_str().unwrap_or("");
                let arguments_str = function["arguments"].as_str().unwrap_or("{}");

                let args: Value = serde_json::from_str(arguments_str).unwrap_or(json!({}));

                slog!(
                    "[openai] Calling MCP tool '{}' for {}:{}",
                    tool_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: {tool_name}]\n"),
                });

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

                let result_content = 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()),
                        });
                        result.clone()
                    }
                    Err(e) => {
                        emit(AgentEvent::Output {
                            story_id: ctx.story_id.clone(),
                            agent_name: ctx.agent_name.clone(),
                            text: format!("[Tool error: {e}]\n"),
                        });
                        format!("Error: {e}")
                    }
                };

                // OpenAI expects tool results as role=tool messages with
                // the matching tool_call_id.
                messages.push(json!({
                    "role": "tool",
                    "tool_call_id": call_id,
                    "content": result_content,
                }));
            }
        }
    }

    fn stop(&self) {
        self.cancelled.stop();
    }

    fn get_status(&self) -> RuntimeStatus {
        self.cancelled.status()
    }
}

// ── Helper functions ─────────────────────────────────────────────────

/// Load MCP tool definitions directly and convert to OpenAI function-calling format.
fn convert_mcp_tools_to_openai() -> Vec<Value> {
    let tools = list_tools();

    let mut openai_tools = Vec::new();

    for tool in &tools {
        let name = tool["name"].as_str().unwrap_or("").to_string();
        let description = tool["description"].as_str().unwrap_or("").to_string();

        if name.is_empty() {
            continue;
        }

        // OpenAI function calling uses JSON Schema natively for parameters,
        // so the MCP inputSchema can be used with minimal cleanup.
        let parameters = convert_mcp_schema_to_openai(tool.get("inputSchema"));

        openai_tools.push(json!({
            "type": "function",
            "function": {
                "name": name,
                "description": description,
                "parameters": parameters.unwrap_or_else(|| json!({"type": "object", "properties": {}})),
            }
        }));
    }

    slog!(
        "[openai] Loaded {} MCP tools as function definitions",
        openai_tools.len()
    );
    openai_tools
}

/// Convert an MCP inputSchema (JSON Schema) to OpenAI-compatible
/// function parameters.
///
/// OpenAI uses JSON Schema natively, so less transformation is needed
/// compared to Gemini. We still strip `$schema` to keep payloads clean.
fn convert_mcp_schema_to_openai(schema: Option<&Value>) -> Option<Value> {
    let schema = schema?;

    let mut result = json!({
        "type": "object",
    });

    if let Some(properties) = schema.get("properties") {
        result["properties"] = clean_schema_properties(properties, false);
    } else {
        result["properties"] = json!({});
    }

    if let Some(required) = schema.get("required") {
        result["required"] = required.clone();
    }

    // OpenAI recommends additionalProperties: false for strict mode.
    result["additionalProperties"] = json!(false);

    Some(result)
}

/// Parse token usage from an OpenAI API response.
fn parse_usage(response: &Value) -> Option<TokenUsage> {
    let usage = response.get("usage")?;
    Some(TokenUsage {
        input_tokens: usage
            .get("prompt_tokens")
            .and_then(|v| v.as_u64())
            .unwrap_or(0),
        output_tokens: usage
            .get("completion_tokens")
            .and_then(|v| v.as_u64())
            .unwrap_or(0),
        cache_creation_input_tokens: 0,
        cache_read_input_tokens: 0,
        // OpenAI API doesn't report cost directly; leave at 0.
        total_cost_usd: 0.0,
    })
}

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

#[cfg(test)]
mod tests {
    use super::super::api_common::test_runtime_context;
    use super::*;

    #[test]
    fn convert_mcp_schema_simple_object() {
        let schema = json!({
            "type": "object",
            "properties": {
                "story_id": {
                    "type": "string",
                    "description": "Story identifier"
                }
            },
            "required": ["story_id"]
        });

        let result = convert_mcp_schema_to_openai(Some(&schema)).unwrap();
        assert_eq!(result["type"], "object");
        assert!(result["properties"]["story_id"].is_object());
        assert_eq!(result["required"][0], "story_id");
        assert_eq!(result["additionalProperties"], false);
    }

    #[test]
    fn convert_mcp_schema_empty_properties() {
        let schema = json!({
            "type": "object",
            "properties": {}
        });

        let result = convert_mcp_schema_to_openai(Some(&schema)).unwrap();
        assert_eq!(result["type"], "object");
        assert!(result["properties"].as_object().unwrap().is_empty());
    }

    #[test]
    fn convert_mcp_schema_none_returns_none() {
        assert!(convert_mcp_schema_to_openai(None).is_none());
    }

    #[test]
    fn convert_mcp_schema_strips_dollar_schema() {
        let schema = json!({
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "$schema": "http://json-schema.org/draft-07/schema#"
                }
            }
        });

        let result = convert_mcp_schema_to_openai(Some(&schema)).unwrap();
        let name_prop = &result["properties"]["name"];
        assert!(name_prop.get("$schema").is_none());
        assert_eq!(name_prop["type"], "string");
    }

    #[test]
    fn parse_usage_valid() {
        let response = json!({
            "usage": {
                "prompt_tokens": 100,
                "completion_tokens": 50,
                "total_tokens": 150
            }
        });

        let usage = parse_usage(&response).unwrap();
        assert_eq!(usage.input_tokens, 100);
        assert_eq!(usage.output_tokens, 50);
        assert_eq!(usage.cache_creation_input_tokens, 0);
        assert_eq!(usage.total_cost_usd, 0.0);
    }

    #[test]
    fn parse_usage_missing() {
        let response = json!({"choices": []});
        assert!(parse_usage(&response).is_none());
    }

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

    #[test]
    fn model_extraction_from_command_gpt() {
        let ctx = test_runtime_context("gpt-4o", vec![]);
        assert!(ctx.command.starts_with("gpt"));
    }

    #[test]
    fn model_extraction_from_command_o3() {
        let ctx = test_runtime_context("o3", vec![]);
        assert!(ctx.command.starts_with("o"));
    }
}