//! The startup self-test: before `loopd` serves anyone, it checks that the server is the one its
//! config describes, that the model's tags tokenize as recorded, 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 bxcore::sanitize::{TAGS, neutralise};

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,
    /// The first tag, by its recorded id, whose token is not in its place in the server's answer
    /// about the tags (or, past the last tag, an extra tag token the server gave).
    TagNotRecorded {
        id: u32,
    },
    /// The server's answer about the neutralised forms holds the tag token `id`.
    NeutralisedIsTag {
        id: u32,
    },
    /// `/tokenize` failed; `what` is "the tags" or "the neutralised tags".
    Tokenize {
        what: &'static str,
        error: InferError,
    },
}

impl SelfTestError {
    /// The runbook pointer that ends the line `loopd` prints for this failure. A tag failure has
    /// its own entry; every other failure shares the self-test's.
    pub fn pointer(&self) -> &'static str {
        match self {
            SelfTestError::TagNotRecorded { .. }
            | SelfTestError::NeutralisedIsTag { .. }
            | SelfTestError::Tokenize { .. } => "see docs/runbook.md#loopd-tags-unexpected",
            SelfTestError::Mismatch { .. }
            | SelfTestError::ToolCall(_)
            | SelfTestError::CacheMiss { .. }
            | SelfTestError::Infer(_)
            | SelfTestError::Hash => "see docs/runbook.md#loopd-selftest-failed",
        }
    }
}

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"),
            SelfTestError::TagNotRecorded { id } => {
                write!(
                    f,
                    "the tag with id {id} does not tokenize as its recorded token"
                )
            }
            SelfTestError::NeutralisedIsTag { id } => {
                write!(f, "a neutralised tag tokenizes to the tag with id {id}")
            }
            SelfTestError::Tokenize { what, error } => {
                write!(f, "could not tokenize {what}: {error}")
            }
        }
    }
}

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 four 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("tags tokenize as recorded");
    tags_tokenize(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 = bxcore::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 (M5b spec, section 8): asked about every tag of `bxcore::sanitize::TAGS` at once, one
/// per line, the server gives each tag's recorded token, in the list's order; asked about every
/// neutralised form, it gives none of those tokens. Each tag is a special token, which the
/// tokenizer splits out before anything else, so the lines between cannot change it. Two requests,
/// so that the check fits in `inferproxy`'s burst of connections. Like check 1 it sends no prompt
/// and uses no slot. A tag missing from the list cannot be found: the server does not list them.
/// An error names a tag by its id, never by its text, which the operator's terminal and log would
/// show whole.
fn tags_tokenize(client: &Client) -> Result<(), SelfTestError> {
    let recorded: Vec<u32> = TAGS.iter().map(|tag| tag.id).collect();
    let tags: Vec<&str> = TAGS.iter().map(|tag| tag.text).collect();
    let got: Vec<u32> = client
        .token_ids(&tags.join("\n"))
        .map_err(|error| SelfTestError::Tokenize {
            what: "the tags",
            error,
        })?
        .into_iter()
        .filter(|id| recorded.contains(id))
        .collect();
    let places = recorded.len().max(got.len());
    if let Some(at) = (0..places).find(|&at| got.get(at) != recorded.get(at)) {
        let id = recorded
            .get(at)
            .or(got.get(at))
            .copied()
            .unwrap_or_default();
        return Err(SelfTestError::TagNotRecorded { id });
    }
    let forms: Vec<String> = TAGS
        .iter()
        .map(|tag| neutralise(tag.text).into_owned())
        .collect();
    let got = client
        .token_ids(&forms.join("\n"))
        .map_err(|error| SelfTestError::Tokenize {
            what: "the neutralised tags",
            error,
        })?;
    match got.into_iter().find(|id| recorded.contains(id)) {
        Some(id) => Err(SelfTestError::NeutralisedIsTag { id }),
        None => Ok(()),
    }
}

/// Check 3: 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 4: 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 }),
    }
}