import Foundation

public enum Gemma4Format {
    public static func buildChatPrompt(messages: [[String: Any]], tools: [[String: Any]]? = nil) -> String {
        var prompt = "<bos>"

        if let tools = tools, !tools.isEmpty {
            let toolDefs = tools.compactMap { t -> String? in
                guard let fn = t["function"] as? [String: Any],
                      let name = fn["name"] as? String else { return nil }
                return name
            }
            if !toolDefs.isEmpty {
                prompt += "<|tool_def|>"
                for (i, name) in toolDefs.enumerated() {
                    if i > 0 { prompt += "|" }
                    prompt += name
                }
                prompt += "<|tool_def|>"
            }
        }

        for msg in messages {
            guard let role = msg["role"] as? String else { continue }

            switch role {
            case "system":
                prompt += "<|turn|>system\n"
                if let content = msg["content"] as? String {
                    prompt += content
                }
                prompt += "<turn|>"

            case "user":
                prompt += "<|turn|>user\n"
                if let content = msg["content"] as? String {
                    prompt += content
                }
                prompt += "<turn|>"

            case "assistant":
                prompt += "<|turn|>model\n"
                if let content = msg["content"] as? String {
                    prompt += content
                }
                if let toolCalls = msg["tool_calls"] as? [[String: Any]] {
                    for tc in toolCalls {
                        if let fn = tc["function"] as? [String: Any],
                           let name = fn["name"] as? String,
                           let args = fn["arguments"] {
                            let argsStr: String
                            if let s = args as? String { argsStr = s }
                            else if let d = try? JSONSerialization.data(withJSONObject: args),
                                    let s = String(data: d, encoding: .utf8) { argsStr = s }
                            else { argsStr = "{}" }
                            prompt += " <|tool_call|>\(name):\(argsStr)<|tool_call|>"
                        }
                    }
                }
                prompt += "<turn|>"

            case "tool":
                prompt += "<|turn|>tool\n"
                if let content = msg["content"] as? String {
                    prompt += content
                }
                if let callId = msg["tool_call_id"] as? String {
                    prompt += " [call_id: \(callId)]"
                }
                prompt += "<turn|>"

            default:
                break
            }
        }

        prompt += "<|turn|>model\n"
        return prompt
    }

    public static func gemma4ArgsToJSON(_ rawArgs: String) -> String {
        let trimmed = rawArgs.trimmingCharacters(in: .whitespacesAndNewlines)
        // Try to parse as JSON first
        if let data = trimmed.data(using: .utf8),
           let _ = try? JSONSerialization.jsonObject(with: data) {
            return trimmed
        }
        // Convert key=value or key:value pairs to JSON
        var entries: [String: String] = [:]
        for part in trimmed.components(separatedBy: ",") {
            let kv = part.split(separator: "=", maxSplits: 1).map(String.init)
            if kv.count == 2 {
                let key = kv[0].trimmingCharacters(in: .whitespaces)
                var val = kv[1].trimmingCharacters(in: .whitespaces)
                // Remove any trailing punctuation
                while val.last == "," || val.last == ")" || val.last == "]" {
                    val = String(val.dropLast())
                }
                entries[key] = val
            }
        }
        if !entries.isEmpty {
            if let data = try? JSONSerialization.data(withJSONObject: entries),
               let json = String(data: data, encoding: .utf8) {
                return json
            }
        }
        return "{\"value\":\"\(trimmed.replacingOccurrences(of: "\"", with: "\\\""))\"}"
    }
}

public struct ToolCallResult: Codable, Sendable {
    public let id: String
    public let function: ToolCallFunction

    public init(id: String, function: ToolCallFunction) {
        self.id = id
        self.function = function
    }
}

public struct ToolCallFunction: Codable, Sendable {
    public let name: String
    public let arguments: String

    public init(name: String, arguments: String) {
        self.name = name
        self.arguments = arguments
    }
}