//
//  ToolRegistry.swift
//  Zyquo Agent
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  The catalog of tools the agent can use. Hands the model the ToolSpecs,
//  resolves incoming ToolCalls by name, decodes and validates their JSON
//  arguments against the tool's schema, executes, and normalizes every
//  failure class (unknown tool, malformed JSON, missing parameters, policy
//  denial, cancellation) into an error ToolResult the model can act on.
//  Adding a tool = conform to `Tool` + append it here (or pass a custom list
//  to `init`).
//

import Foundation

struct ToolRegistry: Sendable {
    private let tools: [any Tool]
    private let toolsByName: [String: any Tool]

    init(tools: [any Tool]) {
        self.tools = tools
        var index: [String: any Tool] = [:]
        for tool in tools {
            precondition(index[tool.name] == nil, "Duplicate tool name: \(tool.name)")
            index[tool.name] = tool
        }
        self.toolsByName = index
    }

    /// The standard Zyquo Agent tool set: bash, osascript, and the five
    /// workspace file tools.
    static func standard(executor: ExecutionService) -> ToolRegistry {
        ToolRegistry(tools: [
            ShellTool(executor: executor),
            AppleScriptTool(executor: executor),
            ReadFileTool(),
            WriteFileTool(),
            EditFileTool(),
            ListDirTool(),
            SearchFilesTool(),
        ])
    }

    /// Provider-neutral specs handed to the model on every request.
    var toolSpecs: [ToolSpec] {
        tools.map(\.toolSpec)
    }

    /// All registered tool names, in registration order.
    var toolNames: [String] {
        tools.map(\.name)
    }

    func tool(named name: String) -> (any Tool)? {
        toolsByName[name]
    }

    /// Executes one model-issued tool call end to end and returns the
    /// ToolResult to thread back into the conversation. Never throws for
    /// tool-level failures — the model must see them as error results so it
    /// can self-correct; only Task cancellation escapes as a thrown error.
    func execute(call: ToolCall, context: ToolExecutionContext) async throws -> ToolResult {
        guard let tool = toolsByName[call.name] else {
            return ToolResult(
                toolCallID: call.id,
                content: "Unknown tool “\(call.name)”. Available tools: \(toolNames.joined(separator: ", ")).",
                isError: true
            )
        }

        // Decode the raw argument JSON. An empty string counts as {} —
        // some providers omit arguments for parameterless calls.
        let rawJSON = call.argumentsJSON.trimmingCharacters(in: .whitespacesAndNewlines)
        let arguments: JSONValue
        if rawJSON.isEmpty {
            arguments = .object([:])
        } else if let parsed = JSONValue.parse(rawJSON) {
            arguments = parsed
        } else {
            return ToolResult(
                toolCallID: call.id,
                content: "The arguments for \(call.name) were not valid JSON. Re-issue the call with a well-formed JSON object.",
                isError: true
            )
        }
        guard case .object = arguments else {
            return ToolResult(
                toolCallID: call.id,
                content: "The arguments for \(call.name) must be a JSON object, got: \(rawJSON.prefix(100)).",
                isError: true
            )
        }

        // Validate against the tool's schema (required keys + basic types).
        if let problem = Self.validate(arguments: arguments, against: tool.parametersSchema, toolName: call.name) {
            return ToolResult(toolCallID: call.id, content: problem, isError: true)
        }

        do {
            let outcome = try await tool.execute(arguments: arguments, context: context)
            return ToolResult(toolCallID: call.id, content: outcome.content, isError: outcome.isError)
        } catch is CancellationError {
            throw CancellationError()
        } catch let denial as PolicyDenied {
            // Tools normally convert denials themselves; this is a backstop.
            return ToolResult(
                toolCallID: call.id,
                content: "Action denied by the safety policy: \(denial.reason)",
                isError: true
            )
        } catch {
            return ToolResult(
                toolCallID: call.id,
                content: "\(call.name) failed: \(error.localizedDescription)",
                isError: true
            )
        }
    }

    // MARK: Schema validation

    /// Checks `required` properties are present and that provided values
    /// match the schema's declared primitive types. Returns a model-actionable
    /// problem description, or nil when valid. (Full JSON-Schema validation
    /// is intentionally out of scope — schemas stay in the flat common
    /// subset per docs/AGENT-RESEARCH.md §2.1.)
    static func validate(arguments: JSONValue, against schema: JSONValue, toolName: String) -> String? {
        guard case .object(let args) = arguments,
              case .object(let schemaObject) = schema else { return nil }

        if case .array(let required)? = schemaObject["required"] {
            for entry in required {
                if case .string(let key) = entry, args[key] == nil {
                    return "\(toolName): missing required parameter `\(key)`."
                }
            }
        }

        if case .object(let properties)? = schemaObject["properties"] {
            for (key, value) in args {
                guard case .object(let property)? = properties[key] else {
                    // Unknown extra keys are tolerated (models add them).
                    continue
                }
                guard case .string(let expected)? = property["type"] else { continue }
                if let problem = typeMismatch(value: value, expected: expected, key: key, toolName: toolName) {
                    return problem
                }
            }
        }
        return nil
    }

    private static func typeMismatch(value: JSONValue, expected: String, key: String, toolName: String) -> String? {
        let actual: String
        switch value {
        case .null: actual = "null"
        case .bool: actual = "boolean"
        case .number(let n): actual = n.truncatingRemainder(dividingBy: 1) == 0 ? "integer" : "number"
        case .string: actual = "string"
        case .array: actual = "array"
        case .object: actual = "object"
        }
        let compatible: Bool
        switch expected {
        case "integer": compatible = actual == "integer"
        case "number": compatible = actual == "integer" || actual == "number"
        case "boolean": compatible = actual == "boolean"
        case "string": compatible = actual == "string"
        case "array": compatible = actual == "array"
        case "object": compatible = actual == "object"
        default: compatible = true
        }
        if !compatible {
            return "\(toolName): parameter `\(key)` should be \(expected), got \(actual)."
        }
        return nil
    }
}