use std::collections::HashMap;
use std::time::Duration;

use ai::diff_validation::ParsedDiff;
use ai::skills::{SkillPathOrigin, SkillReference};
use galaxy_agent_core::ToolCall;
use uuid::Uuid;

use super::rig_request::MCPToolTarget;
use crate::ai::agent::task::TaskId;
use crate::ai::agent::{
    AIAgentAction, AIAgentActionType, AIAgentPtyWriteMode, AskUserQuestionItem,
    AskUserQuestionOption, AskUserQuestionType, CreateDocumentsRequest, DocumentDiff,
    DocumentToCreate, EditDocumentsRequest, FileEdit, FileLocations, ReadDocumentsRequest,
    ReadFilesRequest, ReadSkillRequest, RunAgentsAgentRunConfig, RunAgentsExecutionMode,
    RunAgentsRequest, SearchCodebaseRequest, ShellCommandDelay, StartAgentExecutionMode,
    StartAgentVersion,
};
use crate::ai::document::ai_document_model::AIDocumentId;

pub(crate) fn action_from_tool_call(
    task_id: &str,
    call: &ToolCall,
    skill_path_origin: &SkillPathOrigin,
    mcp_tool_aliases: &HashMap<String, MCPToolTarget>,
) -> Result<AIAgentAction, String> {
    let input = &call.arguments;
    if !input.is_object() {
        return Err(format!("invalid {} input: expected an object", call.name));
    }
    let action = if let Some(target) = mcp_tool_aliases.get(&call.name) {
        AIAgentActionType::CallMCPTool {
            server_id: target.server_id,
            name: target.name.clone(),
            input: input.clone(),
        }
    } else {
        match call.name.as_str() {
            "run_shell_command" => AIAgentActionType::RequestCommandOutput {
                command: required_nonempty_string(input, "command")?,
                is_read_only: Some(optional_boolean(input, "is_read_only")?.unwrap_or(false)),
                is_risky: Some(optional_boolean(input, "is_risky")?.unwrap_or(false)),
                wait_until_completion: optional_boolean(input, "wait_until_complete")?
                    .unwrap_or(false),
                uses_pager: Some(optional_boolean(input, "uses_pager")?.unwrap_or(false)),
                rationale: None,
                citations: Vec::new(),
            },
            "read_files" => AIAgentActionType::ReadFiles(ReadFilesRequest {
                locations: required_array(input, "files")?
                    .iter()
                    .enumerate()
                    .map(|(index, file)| file_location(file, index))
                    .collect::<Result<_, _>>()?,
            }),
            "apply_file_diffs" => AIAgentActionType::RequestFileEdits {
                file_edits: file_edits(input)?,
                title: Some(required_string(input, "summary")?),
            },
            "grep" => AIAgentActionType::Grep {
                queries: required_strings(input, "queries")?,
                path: optional_string(input, "path")?.unwrap_or_default(),
            },
            "file_glob" => AIAgentActionType::FileGlob {
                patterns: required_strings(input, "patterns")?,
                path: optional_string(input, "path")?.filter(|path| !path.is_empty()),
            },
            "search_codebase" => AIAgentActionType::SearchCodebase(SearchCodebaseRequest {
                query: required_string(input, "query")?,
                partial_paths: optional_strings(input, "path_filters")?
                    .filter(|paths| !paths.is_empty()),
                codebase_path: optional_string(input, "path")?.filter(|path| !path.is_empty()),
            }),
            "write_to_long_running_shell_command" => {
                AIAgentActionType::WriteToLongRunningShellCommand {
                    block_id: required_nonempty_string(input, "command_id")?.into(),
                    input: required_string(input, "input")?.into_bytes().into(),
                    mode: match optional_string(input, "mode")?.as_deref() {
                        Some("line") => AIAgentPtyWriteMode::Line,
                        Some("block") => AIAgentPtyWriteMode::Block,
                        Some("raw") | None => AIAgentPtyWriteMode::Raw,
                        Some(mode) => {
                            return Err(format!(
                            "invalid field \"mode\": expected \"raw\", \"line\", or \"block\", got {mode:?}"
                        ));
                        }
                    },
                }
            }
            "interrupt_shell_command" => AIAgentActionType::WriteToLongRunningShellCommand {
                block_id: required_nonempty_string(input, "command_id")?.into(),
                input: vec![galaxy_terminal::model::escape_sequences::C0::ETX].into(),
                mode: AIAgentPtyWriteMode::Raw,
            },
            "read_shell_command_output" => AIAgentActionType::ReadShellCommandOutput {
                block_id: required_nonempty_string(input, "command_id")?.into(),
                delay: Some(ShellCommandDelay::Duration(Duration::from_secs(
                    optional_bounded_u64(
                        input,
                        "wait_seconds",
                        crate::ai::bedrock::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS,
                    )?
                    .unwrap_or(2),
                ))),
            },
            "read_mcp_resource" => AIAgentActionType::ReadMCPResource {
                server_id: Some(required_uuid(input, "server_id")?),
                name: String::new(),
                uri: Some(required_string(input, "uri")?),
            },
            "read_plan" | "read_documents" | "read_notebook" => {
                AIAgentActionType::ReadDocuments(ReadDocumentsRequest {
                    document_ids: required_strings(input, "document_ids")?
                        .into_iter()
                        .map(|id| {
                            AIDocumentId::try_from(id.clone()).map_err(|_| {
                                format!("invalid document_ids entry: {id:?} is not a document ID")
                            })
                        })
                        .collect::<Result<_, _>>()?,
                })
            }
            "create_plan" | "create_documents" | "create_notebook" => {
                AIAgentActionType::CreateDocuments(CreateDocumentsRequest {
                    documents: required_array(input, "documents")?
                        .iter()
                        .enumerate()
                        .map(|(index, document)| {
                            require_object(document, &format!("documents[{index}]"))?;
                            Ok(DocumentToCreate {
                                title: required_string(document, "title")?,
                                content: required_string(document, "content")?,
                            })
                        })
                        .collect::<Result<_, String>>()?,
                })
            }
            "edit_plan" | "edit_documents" | "edit_notebook" => {
                AIAgentActionType::EditDocuments(EditDocumentsRequest {
                    diffs: required_array(input, "diffs")?
                        .iter()
                        .enumerate()
                        .map(|(index, diff)| {
                            require_object(diff, &format!("diffs[{index}]"))?;
                            let document_id = required_string(diff, "document_id")?;
                            Ok(DocumentDiff {
                                document_id: AIDocumentId::try_from(document_id.clone())
                                    .map_err(|_| format!("invalid document_id: {document_id:?}"))?,
                                search: required_string(diff, "search")?,
                                replace: required_string(diff, "replace")?,
                            })
                        })
                        .collect::<Result<_, String>>()?,
                })
            }
            "run_agents" => AIAgentActionType::RunAgents(RunAgentsRequest {
                summary: required_nonempty_string(input, "summary")?,
                base_prompt: optional_string(input, "base_prompt")?.unwrap_or_default(),
                skills: skill_references(input, skill_path_origin)?,
                model_id: optional_string(input, "model_id")?.unwrap_or_default(),
                harness_type: optional_string(input, "harness_type")?.unwrap_or_default(),
                execution_mode: run_agents_execution_mode(input)?,
                agent_run_configs: nonempty_required_array(input, "agent_run_configs")?
                    .iter()
                    .enumerate()
                    .map(|(index, config)| {
                        require_object(config, &format!("agent_run_configs[{index}]"))?;
                        Ok(RunAgentsAgentRunConfig {
                            name: required_nonempty_string(config, "name")?,
                            prompt: required_nonempty_string(config, "prompt")?,
                            title: optional_string(config, "title")?.unwrap_or_default(),
                        })
                    })
                    .collect::<Result<_, String>>()?,
                plan_id: optional_string(input, "plan_id")?.unwrap_or_default(),
                harness_auth_secret_name: None,
            }),
            "start_agent" => AIAgentActionType::StartAgent {
                version: StartAgentVersion::V1,
                name: required_nonempty_string(input, "name")?,
                prompt: required_nonempty_string(input, "prompt")?,
                execution_mode: StartAgentExecutionMode::local_with_defaults(),
                lifecycle_subscription: None,
            },
            "send_message_to_agent" => AIAgentActionType::SendMessageToAgent {
                addresses: vec![required_string(input, "agent_id")?],
                subject: String::new(),
                message: required_string(input, "message")?,
            },
            "transfer_shell_command_control_to_user" => {
                AIAgentActionType::TransferShellCommandControlToUser {
                    reason: required_nonempty_string(input, "reason")?,
                }
            }
            "wait_for_events" => AIAgentActionType::WaitForEvents {
                tool_call_id: call.id.clone(),
                idle_timeout_seconds: optional_nonnegative_i32(input, "idle_timeout_seconds")?
                    .unwrap_or(0),
            },
            "ask_user_question" => AIAgentActionType::AskUserQuestion {
                questions: vec![AskUserQuestionItem {
                    question_id: Uuid::new_v4().to_string(),
                    question: required_string(input, "question")?,
                    question_type: AskUserQuestionType::MultipleChoice {
                        is_multiselect: false,
                        options: optional_strings(input, "options")?
                            .unwrap_or_default()
                            .into_iter()
                            .enumerate()
                            .map(|(index, label)| AskUserQuestionOption {
                                label,
                                recommended: index == 0,
                            })
                            .collect(),
                        supports_other: true,
                    },
                }],
            },
            "read_skill" => {
                let skill = required_string(input, "skill")?;
                let skill = match required_string(input, "reference_type")?.as_str() {
                    "bundled" => SkillReference::BundledSkillId(skill),
                    "path" => SkillReference::Path(
                        skill_path_origin
                            .location_for_path(skill)
                            .map_err(|error| error.to_string())?,
                    ),
                    reference_type => {
                        return Err(format!(
                        "invalid reference_type: expected \"path\" or \"bundled\", got {reference_type:?}"
                    ));
                    }
                };
                AIAgentActionType::ReadSkill(ReadSkillRequest { skill })
            }
            "fetch_conversation" => AIAgentActionType::FetchConversation {
                conversation_id: required_string(input, "conversation_id")?,
            },
            name if name.starts_with("mcp__") => {
                let mut parts = name.splitn(3, "__");
                let _prefix = parts.next();
                let server_id = parts.next().and_then(|value| Uuid::parse_str(value).ok());
                let name = parts
                    .next()
                    .unwrap_or_else(|| name.strip_prefix("mcp__").unwrap_or(name))
                    .to_string();
                AIAgentActionType::CallMCPTool {
                    server_id,
                    name,
                    input: input.clone(),
                }
            }
            name => return Err(format!("unsupported Rig tool proposal: {name}")),
        }
    };

    let tool_name = matches!(
        call.name.as_str(),
        "read_notebook" | "create_notebook" | "edit_notebook"
    )
    .then(|| "notebook".to_string());
    Ok(AIAgentAction {
        id: call.id.clone().into(),
        task_id: TaskId::new(task_id.to_string()),
        action,
        requires_result: true,
        tool_name,
    })
}

fn require_object(input: &serde_json::Value, field: &str) -> Result<(), String> {
    input
        .is_object()
        .then_some(())
        .ok_or_else(|| format!("invalid {field}: expected an object"))
}

fn required_string(input: &serde_json::Value, key: &str) -> Result<String, String> {
    input
        .get(key)
        .ok_or_else(|| format!("missing required field {key:?}"))?
        .as_str()
        .map(ToOwned::to_owned)
        .ok_or_else(|| format!("invalid field {key:?}: expected a string"))
}

fn required_nonempty_string(input: &serde_json::Value, key: &str) -> Result<String, String> {
    let value = required_string(input, key)?;
    if value.trim().is_empty() {
        return Err(format!(
            "invalid field {key:?}: expected a non-empty string"
        ));
    }
    Ok(value)
}

fn optional_string(input: &serde_json::Value, key: &str) -> Result<Option<String>, String> {
    input
        .get(key)
        .map(|value| {
            value
                .as_str()
                .map(ToOwned::to_owned)
                .ok_or_else(|| format!("invalid field {key:?}: expected a string"))
        })
        .transpose()
}

fn required_array<'a>(
    input: &'a serde_json::Value,
    key: &str,
) -> Result<&'a Vec<serde_json::Value>, String> {
    input
        .get(key)
        .ok_or_else(|| format!("missing required field {key:?}"))?
        .as_array()
        .ok_or_else(|| format!("invalid field {key:?}: expected an array"))
}

fn nonempty_required_array<'a>(
    input: &'a serde_json::Value,
    key: &str,
) -> Result<&'a Vec<serde_json::Value>, String> {
    let values = required_array(input, key)?;
    if values.is_empty() {
        return Err(format!("invalid field {key:?}: expected at least one item"));
    }
    Ok(values)
}

fn required_strings(input: &serde_json::Value, key: &str) -> Result<Vec<String>, String> {
    strings_from_array(required_array(input, key)?, key)
}

fn optional_strings(input: &serde_json::Value, key: &str) -> Result<Option<Vec<String>>, String> {
    input
        .get(key)
        .map(|value| {
            let values = value
                .as_array()
                .ok_or_else(|| format!("invalid field {key:?}: expected an array"))?;
            strings_from_array(values, key)
        })
        .transpose()
}

fn strings_from_array(values: &[serde_json::Value], key: &str) -> Result<Vec<String>, String> {
    values
        .iter()
        .enumerate()
        .map(|(index, value)| {
            value
                .as_str()
                .map(ToOwned::to_owned)
                .ok_or_else(|| format!("invalid {key}[{index}]: expected a string"))
        })
        .collect()
}

fn required_uuid(input: &serde_json::Value, key: &str) -> Result<Uuid, String> {
    let value = required_string(input, key)?;
    Uuid::parse_str(&value).map_err(|_| format!("invalid field {key:?}: expected a UUID"))
}

fn optional_boolean(input: &serde_json::Value, key: &str) -> Result<Option<bool>, String> {
    input
        .get(key)
        .map(|value| {
            value
                .as_bool()
                .ok_or_else(|| format!("invalid field {key:?}: expected a boolean"))
        })
        .transpose()
}

fn optional_bounded_u64(
    input: &serde_json::Value,
    key: &str,
    maximum: u64,
) -> Result<Option<u64>, String> {
    input
        .get(key)
        .map(|value| {
            let value = value
                .as_u64()
                .ok_or_else(|| format!("invalid field {key:?}: expected a non-negative integer"))?;
            if value > maximum {
                return Err(format!(
                    "invalid field {key:?}: expected an integer no greater than {maximum}"
                ));
            }
            Ok(value)
        })
        .transpose()
}

fn optional_nonnegative_i32(input: &serde_json::Value, key: &str) -> Result<Option<i32>, String> {
    input
        .get(key)
        .map(|value| {
            value
                .as_i64()
                .and_then(|value| i32::try_from(value).ok())
                .filter(|value| *value >= 0)
                .ok_or_else(|| {
                    format!("invalid field {key:?}: expected a non-negative 32-bit integer")
                })
        })
        .transpose()
}

fn skill_references(
    input: &serde_json::Value,
    skill_path_origin: &SkillPathOrigin,
) -> Result<Vec<SkillReference>, String> {
    let Some(skills) = optional_array(input, "skills")? else {
        return Ok(Vec::new());
    };
    skills
        .iter()
        .enumerate()
        .map(|(index, skill)| {
            require_object(skill, &format!("skills[{index}]"))?;
            let reference = required_string(skill, "skill")?;
            match required_string(skill, "reference_type")?.as_str() {
                "bundled" => Ok(SkillReference::BundledSkillId(reference)),
                "path" => skill_path_origin
                    .location_for_path(reference)
                    .map(SkillReference::Path)
                    .map_err(|error| error.to_string()),
                reference_type => Err(format!(
                    "invalid skills[{index}].reference_type: expected \"path\" or \"bundled\", got {reference_type:?}"
                )),
            }
        })
        .collect()
}

fn run_agents_execution_mode(input: &serde_json::Value) -> Result<RunAgentsExecutionMode, String> {
    let Some(execution_mode) = input.get("execution_mode") else {
        return Ok(RunAgentsExecutionMode::Local);
    };
    require_object(execution_mode, "execution_mode")?;
    match optional_string(execution_mode, "type")?.as_deref() {
        Some("remote") => Ok(RunAgentsExecutionMode::Remote {
            environment_id: optional_string(execution_mode, "environment_id")?.unwrap_or_default(),
            worker_host: optional_string(execution_mode, "worker_host")?.unwrap_or_default(),
            computer_use_enabled: optional_boolean(execution_mode, "computer_use_enabled")?
                .unwrap_or(false),
        }),
        Some("local") | None => {
            optional_string(execution_mode, "environment_id")?;
            optional_string(execution_mode, "worker_host")?;
            optional_boolean(execution_mode, "computer_use_enabled")?;
            Ok(RunAgentsExecutionMode::Local)
        }
        Some(mode_type) => Err(format!(
            "invalid execution_mode.type: expected \"local\" or \"remote\", got {mode_type:?}"
        )),
    }
}

fn file_location(file: &serde_json::Value, file_index: usize) -> Result<FileLocations, String> {
    if let Some(name) = file.as_str() {
        return Ok(FileLocations {
            name: name.to_string(),
            lines: Vec::new(),
        });
    }
    require_object(file, &format!("files[{file_index}]"))?;
    let name = required_string(file, "path")?;
    let lines = match file.get("line_ranges") {
        None => Vec::new(),
        Some(value) => value
            .as_array()
            .ok_or_else(|| format!("invalid files[{file_index}].line_ranges: expected an array"))?
            .iter()
            .enumerate()
            .map(|(range_index, range)| {
                require_object(
                    range,
                    &format!("files[{file_index}].line_ranges[{range_index}]"),
                )?;
                let start = required_line_number(range, "start", file_index, range_index)?;
                let inclusive_end = required_line_number(range, "end", file_index, range_index)?;
                if inclusive_end < start {
                    return Err(format!(
                        "invalid files[{file_index}].line_ranges[{range_index}]: end must be greater than or equal to start"
                    ));
                }
                let exclusive_end = inclusive_end.checked_add(1).ok_or_else(|| {
                    format!(
                        "invalid files[{file_index}].line_ranges[{range_index}].end: inclusive end is too large"
                    )
                })?;
                Ok(start..exclusive_end)
            })
            .collect::<Result<_, String>>()?,
    };
    Ok(FileLocations { name, lines })
}

fn required_line_number(
    range: &serde_json::Value,
    key: &str,
    file_index: usize,
    range_index: usize,
) -> Result<usize, String> {
    let value = range
        .get(key)
        .ok_or_else(|| format!("missing required field {key:?}"))?
        .as_u64()
        .and_then(|value| usize::try_from(value).ok())
        .filter(|value| *value > 0)
        .ok_or_else(|| {
            format!(
                "invalid files[{file_index}].line_ranges[{range_index}].{key}: expected a positive integer"
            )
        })?;
    Ok(value)
}

fn file_edits(input: &serde_json::Value) -> Result<Vec<FileEdit>, String> {
    let mut edits = Vec::new();
    if let Some(diffs) = optional_array(input, "diffs")? {
        for (index, diff) in diffs.iter().enumerate() {
            require_object(diff, &format!("diffs[{index}]"))?;
            edits.push(FileEdit::Edit(ParsedDiff::StrReplaceEdit {
                file: Some(required_string(diff, "file_path")?),
                search: Some(required_string(diff, "search")?),
                replace: Some(required_string(diff, "replace")?),
            }));
        }
    }
    if let Some(files) = optional_array(input, "new_files")? {
        for (index, file) in files.iter().enumerate() {
            require_object(file, &format!("new_files[{index}]"))?;
            edits.push(FileEdit::Create {
                file: Some(required_string(file, "file_path")?),
                content: Some(required_string(file, "content")?),
            });
        }
    }
    if let Some(files) = optional_array(input, "deleted_files")? {
        for (index, file) in files.iter().enumerate() {
            let path = file
                .as_str()
                .ok_or_else(|| format!("invalid deleted_files[{index}]: expected a string"))?;
            edits.push(FileEdit::Delete {
                file: Some(path.to_owned()),
            });
        }
    }
    if edits.is_empty() {
        return Err(
            "invalid file edits: expected at least one diff, new file, or deleted file".to_string(),
        );
    }
    Ok(edits)
}

fn optional_array<'a>(
    input: &'a serde_json::Value,
    key: &str,
) -> Result<Option<&'a Vec<serde_json::Value>>, String> {
    input
        .get(key)
        .map(|value| {
            value
                .as_array()
                .ok_or_else(|| format!("invalid field {key:?}: expected an array"))
        })
        .transpose()
}

#[cfg(test)]
#[path = "rig_tool_tests.rs"]
mod tests;