//
//  OpenAINormalizer.swift
//  Zyquo Router
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Canonical OpenAI wire format served by the router. Phase 2 carries the
//  error body and the /v1/models list shapes; Phase 3 adds the full
//  chat/completions request/response/chunk types and normalization.
//

import Foundation
import NIOHTTP1

// MARK: - Error body (spec: {"error": {"message", "type", "param", "code"}})

struct OpenAIError: Codable {
    struct Body: Codable {
        var message: String
        var type: String
        var param: String?
        var code: String?
    }

    var error: Body

    private static let encoder: JSONEncoder = {
        let encoder = JSONEncoder()
        encoder.outputFormatting = [.withoutEscapingSlashes, .sortedKeys]
        return encoder
    }()

    /// A complete error response in OpenAI wire shape.
    static func response(
        status: HTTPResponseStatus,
        message: String,
        type: String,
        param: String? = nil,
        code: String? = nil,
        extraHeaders: [(String, String)] = []
    ) -> RouteResult {
        let payload = OpenAIError(error: Body(message: message, type: type, param: param, code: code))
        let body = (try? encoder.encode(payload)) ?? Data(#"{"error":{"message":"internal error"}}"#.utf8)
        return .complete(
            status: status,
            headers: [("Content-Type", "application/json")] + extraHeaders,
            body: body
        )
    }

    /// The exact 404 OpenAI returns for unknown models.
    static func modelNotFound(_ model: String) -> RouteResult {
        response(
            status: .notFound,
            message: "The model `\(model)` does not exist or you do not have access to it.",
            type: "invalid_request_error",
            param: "model",
            code: "model_not_found"
        )
    }
}

// MARK: - /v1/models list shapes

struct OpenAIModelEntry: Codable {
    var id: String
    var object = "model"
    var created: Int
    var ownedBy: String
    /// Router extension: catalog metadata (context, pricing, capabilities).
    var xZyquo: XZyquoModelInfo?

    enum CodingKeys: String, CodingKey {
        case id, object, created
        case ownedBy = "owned_by"
        case xZyquo = "x_zyquo"
    }
}

struct XZyquoModelInfo: Codable {
    var displayName: String
    var contextWindow: Int
    var maxOutputTokens: Int?
    var vision: Bool
    var tools: Bool
    var reasoning: Bool
    var inputPerMTok: Double?
    var outputPerMTok: Double?
    var alias: String?

    enum CodingKeys: String, CodingKey {
        case displayName = "display_name"
        case contextWindow = "context_window"
        case maxOutputTokens = "max_output_tokens"
        case vision, tools, reasoning, alias
        case inputPerMTok = "input_per_mtok"
        case outputPerMTok = "output_per_mtok"
    }
}

struct OpenAIModelList: Codable {
    var object = "list"
    var data: [OpenAIModelEntry]
}

// MARK: - Response / chunk construction

/// Builds spec-exact `chat.completion` and `chat.completion.chunk` JSON.
/// One emitter per request: `id` and `created` stay constant for a stream.
struct ChunkEmitter {
    let id: String
    let created: Int
    /// Namespaced router model id, echoed on every chunk/response.
    let model: String

    init(model: String) {
        var hex = ""
        for _ in 0..<12 { hex += String(format: "%x", Int.random(in: 0...15)) }
        id = "chatcmpl-\(hex)"
        created = Int(Date().timeIntervalSince1970)
        self.model = model
    }

    static func serialize(_ object: [String: Any]) -> Data {
        (try? JSONSerialization.data(withJSONObject: object)) ?? Data("{}".utf8)
    }

    private func envelope(delta: [String: Any]?, finishReason: String?) -> [String: Any] {
        [
            "id": id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [[
                "index": 0,
                "delta": delta ?? [:],
                "finish_reason": finishReason as Any,
            ] as [String: Any]],
        ]
    }

    /// First chunk of every stream: the role delta.
    func roleChunk() -> Data {
        Self.serialize(envelope(delta: ["role": "assistant", "content": ""], finishReason: nil))
    }

    func contentChunk(_ text: String) -> Data {
        Self.serialize(envelope(delta: ["content": text], finishReason: nil))
    }

    func reasoningChunk(_ text: String) -> Data {
        Self.serialize(envelope(delta: ["reasoning_content": text], finishReason: nil))
    }

    /// Announces a tool call: id + name once, empty arguments accumulator.
    func toolCallStartChunk(toolIndex: Int, callID: String, name: String) -> Data {
        Self.serialize(envelope(
            delta: ["tool_calls": [[
                "index": toolIndex,
                "id": callID,
                "type": "function",
                "function": ["name": name, "arguments": ""],
            ] as [String: Any]]],
            finishReason: nil
        ))
    }

    func toolCallArgumentsChunk(toolIndex: Int, fragment: String) -> Data {
        Self.serialize(envelope(
            delta: ["tool_calls": [[
                "index": toolIndex,
                "function": ["arguments": fragment],
            ] as [String: Any]]],
            finishReason: nil
        ))
    }

    func finishChunk(reason: String) -> Data {
        Self.serialize(envelope(delta: [:], finishReason: reason))
    }

    /// Usage chunk: empty `choices` array, only when the client asked for it.
    func usageChunk(_ usage: [String: Any]) -> Data {
        Self.serialize([
            "id": id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [] as [Any],
            "usage": usage,
        ])
    }

    /// Complete non-streaming `chat.completion` object.
    func completion(
        message: [String: Any],
        finishReason: String,
        usage: [String: Any],
        extras: [String: Any] = [:]
    ) -> [String: Any] {
        var object: [String: Any] = [
            "id": id,
            "object": "chat.completion",
            "created": created,
            "model": model,
            "choices": [[
                "index": 0,
                "message": message,
                "finish_reason": finishReason,
            ] as [String: Any]],
            "usage": usage,
        ]
        for (key, value) in extras {
            object[key] = value
        }
        return object
    }
}

/// OpenAI-shape usage dictionaries from normalized numbers.
enum UsageBuilder {
    static func build(
        promptTokens: Int,
        completionTokens: Int,
        cachedTokens: Int? = nil,
        reasoningTokens: Int? = nil,
        estimated: Bool = false
    ) -> [String: Any] {
        var usage: [String: Any] = [
            "prompt_tokens": promptTokens,
            "completion_tokens": completionTokens,
            "total_tokens": promptTokens + completionTokens,
        ]
        if let cachedTokens, cachedTokens > 0 {
            usage["prompt_tokens_details"] = ["cached_tokens": cachedTokens]
        }
        if let reasoningTokens, reasoningTokens > 0 {
            usage["completion_tokens_details"] = ["reasoning_tokens": reasoningTokens]
        }
        if estimated {
            usage["x_zyquo"] = ["usage_estimated": true]
        }
        return usage
    }

    /// Rough local estimation (~4 chars/token) used only when the upstream
    /// reports nothing; always flagged via x_zyquo.usage_estimated.
    static func estimateTokens(_ text: String) -> Int {
        max(1, text.count / 4)
    }
}