use super::{GlobalConfig, Input, RoleLike, SUBQUERY_COMPRESSION_PROMPT};

use crate::client::{call_chat_completions, call_chat_completions_streaming, Message, MessageRole};
use crate::function::ToolResult;
use crate::utils::{
    create_abort_signal, current_query_scope, dimmed_text, format_query_path, indent_text,
    next_query_id, with_query_scope, AbortSignal, IS_STDOUT_TERMINAL,
};

use anyhow::{anyhow, bail, Result};
use serde_json::{json, Value};

pub const SUBQUERY_TOOL_NAME: &str = "subquery";

pub fn subquery_declaration() -> crate::function::FunctionDeclaration {
    crate::function::FunctionDeclaration {
        name: SUBQUERY_TOOL_NAME.to_string(),
        description: concat!(
            "Delegate a fully self-contained sub task to a fresh LLM query that does NOT ",
            "inherit the current conversation context. Only the context explicitly passed ",
            "through the parameters (by message indexes and/or text excerpts) is visible to ",
            "the subquery. The subquery itself can still use tools (including nested ",
            "subqueries when the depth limit allows), and its final answer is returned to you."
        )
        .to_string(),
        parameters: crate::function::JsonSchema {
            type_value: Some("object".to_string()),
            description: None,
            properties: Some(
                [
                    (
                        "task",
                        crate::function::JsonSchema {
                            type_value: Some("string".to_string()),
                            description: Some(
                                "The task or question for the subquery. Must be fully self-contained; do not rely on the surrounding conversation."
                                .to_string(),
                            ),
                            properties: None,
                            items: None,
                            any_of: None,
                            enum_value: None,
                            default: None,
                            required: None,
                        },
                    ),
                    (
                        "message_indexes",
                        crate::function::JsonSchema {
                            type_value: Some("array".to_string()),
                            description: Some(
                                "0-based indexes of the current context window messages to pass to the subquery as context. Messages are numbered in the order they appear in the context window, starting at 0."
                                .to_string(),
                            ),
                            properties: None,
                            items: Some(Box::new(crate::function::JsonSchema {
                                        type_value: Some("integer".to_string()),
                                        description: None,
                                        properties: None,
                                        items: None,
                                        any_of: None,
                                        enum_value: None,
                                        default: None,
                                        required: None,
                            })),
                            any_of: None,
                            enum_value: None,
                            default: None,
                            required: None,
                        },
                    ),
                    (
                        "excerpts",
                        crate::function::JsonSchema {
                            type_value: Some("array".to_string()),
                            description: Some(
                                "Verbatim text excerpts copied from the current context window, provided to the subquery as extra context."
                                .to_string(),
                            ),
                            properties: None,
                            items: Some(Box::new(crate::function::JsonSchema {
                                        type_value: Some("string".to_string()),
                                        description: None,
                                        properties: None,
                                        items: None,
                                        any_of: None,
                                        enum_value: None,
                                        default: None,
                                        required: None,
                            })),
                            any_of: None,
                            enum_value: None,
                            default: None,
                            required: None,
                        },
                    ),
                ]
                .into_iter()
                .map(|(k, v)| (k.to_string(), v))
                .collect(),
            ),
            items: None,
            any_of: None,
            enum_value: None,
            default: None,
            required: Some(vec!["task".to_string()]),
        },
        agent: false,
    }
}

pub struct SubqueryContext {
    path: Vec<u64>,
    context_messages: Vec<Message>,
}

impl SubqueryContext {
    pub fn from_input(input: &Input) -> Self {
        let path = input.query_path();
        let context_messages = if path.is_empty() {
            vec![]
        } else {
            match input.build_messages() {
                Ok(v) => v,
                Err(err) => {
                    debug!("Failed to build context messages for subquery: {err}");
                    vec![]
                }
            }
        };
        Self {
            path,
            context_messages,
        }
    }
}

pub async fn eval_subquery(
    config: &GlobalConfig,
    ctx: &SubqueryContext,
    arguments: Value,
) -> Result<Value> {
    let parent_path = ctx.path.clone();
    if parent_path.is_empty() {
        bail!("Invalid subquery call: missing parent query id");
    }
    let max_depth = config.read().subquery_max_depth;
    if parent_path.len() >= max_depth {
        bail!("Subquery maximum depth ({max_depth}) exceeded");
    }
    let (task, indexes, excerpts) = parse_subquery_args(&arguments)?;
    if *IS_STDOUT_TERMINAL {
        let arguments_text =
            serde_json::to_string(&arguments).unwrap_or_else(|_| arguments.to_string());
        let line = format!("Call {SUBQUERY_TOOL_NAME} {arguments_text}");
        let line = match current_query_scope() {
            Some(scope) => format!("{scope} {line}"),
            None => line,
        };
        println!("{}", dimmed_text(&line));
    }
    let abort_signal = create_abort_signal();
    let sub_id = next_query_id(&parent_path);
    let mut sub_path = parent_path;
    sub_path.push(sub_id);
    let result = with_query_scope(format_query_path(&sub_path), async {
        run_subquery_inner(config, ctx, task, indexes, excerpts, sub_path, abort_signal).await
    })
    .await;
    match result {
        Ok(text) => {
            let text = maybe_compress_subquery_output(config, text).await;
            Ok(json!({ "output": text }))
        }
        Err(err) => Ok(json!({ "error": format!("{err:?}") })),
    }
}

async fn run_subquery_inner(
    config: &GlobalConfig,
    ctx: &SubqueryContext,
    task: String,
    indexes: Vec<usize>,
    excerpts: Vec<String>,
    sub_path: Vec<u64>,
    abort_signal: AbortSignal,
) -> Result<String> {
    let mut context_parts = vec![];
    if !indexes.is_empty() {
        context_parts.push(format!(
                "The following messages were numerically selected from the parent context window:\n\n{}",
                numbered_context_text(&ctx.context_messages, &indexes)
        ));
    }
    if !excerpts.is_empty() {
        context_parts.push(format!(
            "The following text excerpts were selected from the parent context window:\n\n{}",
            excerpts.join("\n\n")
        ));
    }
    let mut messages = vec![];
    if !context_parts.is_empty() {
        messages.push(Message::new(
            MessageRole::System,
            crate::client::MessageContent::Text(context_parts.join("\n\n")),
        ));
    }
    messages.push(Message::new(
        MessageRole::User,
        crate::client::MessageContent::Text(task),
    ));
    run_subquery(config, messages, sub_path, abort_signal).await
}

fn parse_subquery_args(arguments: &Value) -> Result<(String, Vec<usize>, Vec<String>)> {
    let arguments = match arguments {
        Value::Object(map) => Value::Object(map.clone()),
        Value::String(text) => serde_json::from_str(text)
            .map_err(|_| anyhow!("Invalid arguments for '{SUBQUERY_TOOL_NAME}': {text}"))?,
        _ => bail!("Invalid arguments for '{SUBQUERY_TOOL_NAME}': {arguments}"),
    };
    let task = arguments["task"]
        .as_str()
        .map(|v| v.trim().to_string())
        .unwrap_or_default();
    if task.is_empty() {
        bail!("The call '{SUBQUERY_TOOL_NAME}' requires a non-empty 'task' argument");
    }
    let indexes = match &arguments["message_indexes"] {
        Value::Array(items) => items
            .iter()
            .filter_map(|v| match v {
                Value::String(s) => s.trim().parse::<usize>().ok(),
                _ => v.as_u64().map(|v| v as usize),
            })
            .collect(),
        Value::Null => vec![],
        other => {
            bail!("The call '{SUBQUERY_TOOL_NAME}' has invalid 'message_indexes': {other}")
        }
    };
    let excerpts = match &arguments["excerpts"] {
        Value::Array(items) => items
            .iter()
            .filter_map(|v| v.as_str().map(|v| v.to_string()))
            .collect(),
        Value::String(v) => vec![v.clone()],
        Value::Null => vec![],
        other => bail!("The call '{SUBQUERY_TOOL_NAME}' has invalid 'excerpts': {other}"),
    };
    Ok((task, indexes, excerpts))
}

fn numbered_context_text(messages: &[Message], indexes: &[usize]) -> String {
    let mut lines = vec![];
    for (i, message) in messages.iter().enumerate() {
        if !indexes.contains(&i) {
            continue;
        }
        lines.push(format!(
            "# {i} {}:\n{}",
            format!("{:?}", message.role).to_lowercase(),
            indent_text(message.content.to_text(), 4).trim_start()
        ));
    }
    if lines.is_empty() {
        "(no matching messages found)".into()
    } else {
        lines.join("\n\n")
    }
}

async fn run_subquery(
    config: &GlobalConfig,
    mut messages: Vec<Message>,
    path: Vec<u64>,
    abort_signal: AbortSignal,
) -> Result<String> {
    let client = crate::client::init_client(config, None)?;
    let (temperature, top_p, functions, stream) = {
        let temp_input = Input::from_str(config, "", None);
        let (temperature, top_p) = (temp_input.role().temperature(), temp_input.role().top_p());
        let max_depth = config.read().subquery_max_depth;

        let mut declarations = config
            .read()
            .select_functions(temp_input.role())
            .unwrap_or_default();

        // Never inherit a stale/duplicated subquery entry from select_functions().
        declarations.retain(|v| v.name != SUBQUERY_TOOL_NAME);

        // Only offer nesting while the depth limit allows it.
        if path.len() < max_depth {
            declarations.push(subquery_declaration());
        }

        let functions = if declarations.is_empty() {
            None
        } else {
            Some(declarations)
        };

        (
            temperature,
            top_p,
            functions,
            config.read().stream && !temp_input.role().model().no_stream(),
        )
    };
    let output;
    let prefix = format_query_path(&path);
    loop {
        let input = build_subquery_input(
            config,
            messages.clone(),
            temperature,
            top_p,
            functions.clone(),
            path.clone(),
        )?;
        let (text, tool_results) = if stream {
            call_chat_completions_streaming(
                &input,
                client.as_ref(),
                abort_signal.clone(),
                Some(prefix.clone()),
            )
            .await?
        } else {
            let (text, tool_results) =
                call_chat_completions(&input, false, false, client.as_ref(), abort_signal.clone())
                    .await?;
            print_subquery_output(config, &prefix, &text);
            (text, tool_results)
        };
        if tool_results.is_empty() {
            output = text;
            break;
        }
        messages.extend(tool_result_messages(&text, &tool_results));
    }
    Ok(output)
}

fn build_subquery_input(
    config: &GlobalConfig,
    messages: Vec<Message>,
    temperature: Option<f64>,
    top_p: Option<f64>,
    functions: Option<Vec<crate::function::FunctionDeclaration>>,
    path: Vec<u64>,
) -> Result<Input> {
    let mut input = Input::from_str(config, "", None);
    input.set_subquery(messages, temperature, top_p, functions);
    input.set_query_path(path);
    Ok(input)
}

fn print_subquery_output(config: &GlobalConfig, prefix: &str, text: &str) {
    let text = text.trim();
    if text.is_empty() {
        return;
    }
    if !prefix.is_empty() {
        println!("{}", dimmed_text(prefix));
    }
    let _ = config.read().print_markdown(text);
}

fn tool_result_messages(text: &str, tool_results: &[ToolResult]) -> Vec<Message> {
    let mut messages = vec![];
    let text = text.trim();
    if !text.is_empty() {
        messages.push(Message::new(
            MessageRole::Assistant,
            crate::client::MessageContent::Text(text.to_string()),
        ));
    }
    let results: Vec<ToolResult> = tool_results
        .iter()
        .cloned()
        .enumerate()
        .map(|(i, mut result)| {
            if result.call.id.is_none() {
                result.call.id = Some(format!("subquery-call-{}", i + 1));
            }
            result
        })
        .collect();
    messages.push(Message::new(
        MessageRole::Assistant,
        crate::client::MessageContent::ToolCalls(crate::client::MessageContentToolCalls::new(
            results,
            String::new(),
        )),
    ));
    messages
}

async fn maybe_compress_subquery_output(config: &GlobalConfig, output: String) -> String {
    if !config.read().subquery_compression {
        return output;
    }
    let threshold = config.read().subquery_compression_threshold;
    if threshold == 0 || count_tokens(config, &output) <= threshold {
        return output;
    }
    match compress_subquery_output(config, &output).await {
        Ok(compressed) if !compressed.trim().is_empty() => compressed,
        _ => output,
    }
}

fn count_tokens(config: &GlobalConfig, text: &str) -> usize {
    let input = Input::from_str(config, "", None);
    let messages = [Message::new(
        MessageRole::User,
        crate::client::MessageContent::Text(text.to_string()),
    )];
    input.role().model().total_tokens(&messages)
}

async fn compress_subquery_output(config: &GlobalConfig, output: &str) -> Result<String> {
    let prompt = config
        .read()
        .subquery_compression_prompt
        .clone()
        .unwrap_or_else(|| SUBQUERY_COMPRESSION_PROMPT.to_string());
    let role = config.read().extract_role();
    let text = format!("{prompt}\n\n{output}");
    let input = Input::from_str(config, &text, Some(role));
    input.fetch_chat_text().await
}