//! The startup self-test: before `loopd` serves anyone, it checks that the server is the one its
//! config describes, that a tool call comes back parsed, and that a second turn reuses the first
//! turn's cache. Any check failing means `loopd` does not start.

use crate::llama::{
    ChatMessage, ChatRequest, Client, FinishReason, InferError, ToolSchema,
    info::{CacheOutcome, cache_outcome},
};

#[derive(Debug)]
pub enum SelfTestError {
    Mismatch {
        what: &'static str,
        expected: String,
        got: String,
    },
    ToolCall(String),
    CacheMiss {
        expected: u64,
        got: u64,
    },
    Infer(InferError),
    Hash,
}

impl std::fmt::Display for SelfTestError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SelfTestError::Mismatch {
                what,
                expected,
                got,
            } => write!(f, "{what}: expected {expected}, got {got}"),
            SelfTestError::ToolCall(reason) => {
                write!(f, "the tool call did not come back parsed: {reason}")
            }
            SelfTestError::CacheMiss { expected, got } => write!(
                f,
                "turn 2 did not reuse turn 1's cache: expected {expected}, saw {got}"
            ),
            SelfTestError::Infer(e) => write!(f, "{e}"),
            SelfTestError::Hash => write!(f, "the chat template could not be hashed"),
        }
    }
}

impl std::error::Error for SelfTestError {}

impl From<InferError> for SelfTestError {
    fn from(e: InferError) -> Self {
        SelfTestError::Infer(e)
    }
}

const SYSTEM: &str = "You are Boxmaker, a careful personal agent.";

/// Runs the three checks in order. `on_step` is told the name of each check as it starts.
pub fn run(client: &Client, on_step: &mut dyn FnMut(&str)) -> Result<(), SelfTestError> {
    on_step("server matches config");
    server_matches(client)?;

    on_step("tool call round trip");
    tool_call_round_trip(client)?;

    on_step("turn 2 cache hit");
    cache_hit(client)?;

    Ok(())
}

/// Check 1: the server's chat template hash and slot settings are the ones the config expects.
/// A server that is not the expected one is sent no prompt at all.
fn server_matches(client: &Client) -> Result<(), SelfTestError> {
    let props = client.props().map_err(SelfTestError::Infer)?;
    let expect = &client.config().expect;

    let expected = expect.template_sha256;
    let got = proto::sha256(props.chat_template.as_bytes()).map_err(|_| SelfTestError::Hash)?;
    if expected != got {
        return Err(SelfTestError::Mismatch {
            what: "chat template sha256",
            expected: expected.to_hex(),
            got: got.to_hex(),
        });
    }

    if expect.n_ctx != props.n_ctx {
        return Err(SelfTestError::Mismatch {
            what: "context per slot",
            expected: expect.n_ctx.to_string(),
            got: props.n_ctx.to_string(),
        });
    }

    if expect.slots != props.total_slots {
        return Err(SelfTestError::Mismatch {
            what: "slot count",
            expected: expect.slots.to_string(),
            got: props.total_slots.to_string(),
        });
    }

    Ok(())
}

/// Check 2: a tool call comes back parsed, with a `read_file` tool and a string `path`.
fn tool_call_round_trip(client: &Client) -> Result<(), SelfTestError> {
    let req = ChatRequest {
        slot: client.config().slots.main,
        messages: vec![
            ChatMessage::System {
                content: SYSTEM.to_string(),
            },
            ChatMessage::User {
                content: "Read /etc/hostname and tell me what it says.".to_string(),
            },
        ],
        tools: vec![ToolSchema {
            name: "read_file".to_string(),
            description: "Read a text file and return its contents.".to_string(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": { "path": { "type": "string", "description": "Absolute path" } },
                "required": ["path"],
            }),
        }],
        thinking: false,
    };

    let completion = client
        .chat_with_retry(&req, &mut |_| {})
        .map_err(SelfTestError::Infer)?;

    if completion.finish_reason != FinishReason::ToolCalls {
        return Err(SelfTestError::ToolCall(
            "the completion did not end with a tool call".to_string(),
        ));
    }
    let Some(call) = completion.tool_calls.first() else {
        return Err(SelfTestError::ToolCall(
            "the completion made no tool call".to_string(),
        ));
    };
    if call.name != "read_file" {
        return Err(SelfTestError::ToolCall(format!(
            "the first tool call was {}, not read_file",
            call.name
        )));
    }
    let arguments: serde_json::Value = serde_json::from_str(&call.arguments).map_err(|_| {
        SelfTestError::ToolCall("the tool call arguments were not valid JSON".to_string())
    })?;
    if arguments.get("path").and_then(|v| v.as_str()).is_none() {
        return Err(SelfTestError::ToolCall(
            "the tool call arguments had no string path".to_string(),
        ));
    }

    Ok(())
}

/// Check 3: a second turn that extends the first reuses the first turn's cache.
fn cache_hit(client: &Client) -> Result<(), SelfTestError> {
    let first = ChatRequest {
        slot: client.config().slots.main,
        messages: vec![
            ChatMessage::System {
                content: SYSTEM.to_string(),
            },
            ChatMessage::User {
                content: "Name one colour. One word.".to_string(),
            },
        ],
        tools: Vec::new(),
        thinking: false,
    };
    let turn1 = client
        .chat_with_retry(&first, &mut |_| {})
        .map_err(SelfTestError::Infer)?;

    let mut messages = vec![
        ChatMessage::System {
            content: SYSTEM.to_string(),
        },
        ChatMessage::User {
            content: "Name one colour. One word.".to_string(),
        },
    ];
    messages.push(ChatMessage::Assistant {
        content: turn1.content,
        reasoning_content: turn1.reasoning_content,
        tool_calls: turn1.tool_calls,
    });
    messages.push(ChatMessage::User {
        content: "Name another. One word.".to_string(),
    });
    let second = ChatRequest {
        slot: client.config().slots.main,
        messages,
        tools: Vec::new(),
        thinking: false,
    };
    let turn2 = client
        .chat_with_retry(&second, &mut |_| {})
        .map_err(SelfTestError::Infer)?;

    match cache_outcome(&turn1.timings, &turn2.timings) {
        CacheOutcome::Hit => Ok(()),
        CacheOutcome::Loss { expected, got } => Err(SelfTestError::CacheMiss { expected, got }),
    }
}