import Foundation

/// Model architecture configuration parsed from config.json.
/// Uses JSONSerialization (instead of Codable) to tolerate unknown keys.
public struct ModelConfig: Sendable {
    public let modelType: String?
    public let hiddenSize: Int?
    public let intermediateSize: Int?
    public let numAttentionHeads: Int?
    public let numHiddenLayers: Int?
    public let numKeyValueHeads: Int?
    public let vocabSize: Int?
    public let maxPositionEmbeddings: Int?
    public let rmsNormEps: Float?
    public let ropeTheta: Float?

    // Gemma 4 specific
    public let slidingWindow: Int?
    public let headDim: Int?
    public let globalHeadDim: Int?
    public let slidingHeadDim: Int?
    public let numKvSharedLayers: Int?
    public let hiddenSizePerLayerInput: Int?
    public let slidingWindowPattern: Int?
    public let finalLogitSoftcapping: Float?
    public let tieWordEmbeddings: Bool?
    public let perLayerInputScale: Float?
    public let perLayerProjectionScale: Float?
    public let embedScale: Float?
    /// Per-layer attention type: "full_attention" or "sliding_attention"
    public let layerTypes: [String]?

    // Global KV heads (for full attention layers)
    public let numGlobalKeyValueHeads: Int?

    // K=V sharing (Gemma 4 full attention layers)
    public let attentionKEqualsV: Bool?

    // MoE
    public let enableMoEBlock: Bool?
    public let numExperts: Int?
    public let topKExperts: Int?
    public let moeIntermediateSize: Int?

    public init(
        modelType: String? = nil,
        hiddenSize: Int? = nil,
        intermediateSize: Int? = nil,
        numAttentionHeads: Int? = nil,
        numHiddenLayers: Int? = nil,
        numKeyValueHeads: Int? = nil,
        vocabSize: Int? = nil,
        maxPositionEmbeddings: Int? = nil,
        rmsNormEps: Float? = nil,
        ropeTheta: Float? = nil,
        slidingWindow: Int? = nil,
        headDim: Int? = nil,
        globalHeadDim: Int? = nil,
        slidingHeadDim: Int? = nil,
        numKvSharedLayers: Int? = nil,
        hiddenSizePerLayerInput: Int? = nil,
        slidingWindowPattern: Int? = nil,
        finalLogitSoftcapping: Float? = nil,
        tieWordEmbeddings: Bool? = nil,
        perLayerInputScale: Float? = nil,
        perLayerProjectionScale: Float? = nil,
        embedScale: Float? = nil,
        layerTypes: [String]? = nil,
        numGlobalKeyValueHeads: Int? = nil,
        enableMoEBlock: Bool? = nil,
        numExperts: Int? = nil,
        topKExperts: Int? = nil,
        moeIntermediateSize: Int? = nil,
        attentionKEqualsV: Bool? = nil
    ) {
        self.modelType = modelType
        self.hiddenSize = hiddenSize
        self.intermediateSize = intermediateSize
        self.numAttentionHeads = numAttentionHeads
        self.numHiddenLayers = numHiddenLayers
        self.numKeyValueHeads = numKeyValueHeads
        self.vocabSize = vocabSize
        self.maxPositionEmbeddings = maxPositionEmbeddings
        self.rmsNormEps = rmsNormEps
        self.ropeTheta = ropeTheta
        self.slidingWindow = slidingWindow
        self.headDim = headDim
        self.globalHeadDim = globalHeadDim
        self.slidingHeadDim = slidingHeadDim
        self.numKvSharedLayers = numKvSharedLayers
        self.hiddenSizePerLayerInput = hiddenSizePerLayerInput
        self.slidingWindowPattern = slidingWindowPattern
        self.layerTypes = layerTypes
        self.finalLogitSoftcapping = finalLogitSoftcapping
        self.tieWordEmbeddings = tieWordEmbeddings
        self.perLayerInputScale = perLayerInputScale
        self.perLayerProjectionScale = perLayerProjectionScale
        self.embedScale = embedScale
        self.numGlobalKeyValueHeads = numGlobalKeyValueHeads
        self.enableMoEBlock = enableMoEBlock
        self.numExperts = numExperts
        self.topKExperts = topKExperts
        self.moeIntermediateSize = moeIntermediateSize
        self.attentionKEqualsV = attentionKEqualsV
    }

    /// Load config from a model directory (config.json).
    /// Uses JSONSerialization to tolerate extra keys.
    public static func load(from directory: String) throws -> ModelConfig {
        let url = URL(fileURLWithPath: directory).appendingPathComponent("config.json")
        let data = try Data(contentsOf: url)
        guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            throw WeightError.invalidHeader("config.json is not a dictionary")
        }

        // Some configs nest text params in "text_config"
        let tc = json["text_config"] as? [String: Any] ?? [:]

        return ModelConfig(
            modelType: json.string("model_type"),
            hiddenSize: json.int("hidden_size") ?? tc.int("hidden_size"),
            intermediateSize: json.int("intermediate_size") ?? tc.int("intermediate_size"),
            numAttentionHeads: json.int("num_attention_heads") ?? tc.int("num_attention_heads"),
            numHiddenLayers: json.int("num_hidden_layers") ?? tc.int("num_hidden_layers"),
            numKeyValueHeads: json.int("num_key_value_heads") ?? tc.int("num_key_value_heads"),
            vocabSize: json.int("vocab_size") ?? tc.int("vocab_size"),
            maxPositionEmbeddings: json.int("max_position_embeddings") ?? tc.int("max_position_embeddings"),
            rmsNormEps: json.float("rms_norm_eps") ?? tc.float("rms_norm_eps"),
            ropeTheta: json.float("rope_theta") ?? tc.float("rope_theta"),
            slidingWindow: json.int("sliding_window") ?? tc.int("sliding_window"),
            headDim: json.int("head_dim") ?? tc.int("head_dim"),
            globalHeadDim: json.int("global_head_dim") ?? tc.int("global_head_dim"),
            slidingHeadDim: json.int("sliding_head_dim") ?? tc.int("sliding_head_dim"),
            numKvSharedLayers: json.int("num_kv_shared_layers") ?? tc.int("num_kv_shared_layers"),
            hiddenSizePerLayerInput: json.int("hidden_size_per_layer_input") ?? tc.int("hidden_size_per_layer_input"),
            slidingWindowPattern: json.int("sliding_window_pattern") ?? tc.int("sliding_window_pattern"),
            finalLogitSoftcapping: json.float("final_logit_softcapping") ?? tc.float("final_logit_softcapping"),
            tieWordEmbeddings: json.bool("tie_word_embeddings") ?? tc.bool("tie_word_embeddings"),
            perLayerInputScale: json.float("per_layer_input_scale") ?? tc.float("per_layer_input_scale"),
            perLayerProjectionScale: json.float("per_layer_projection_scale") ?? tc.float("per_layer_projection_scale"),
            embedScale: json.float("embed_scale") ?? tc.float("embed_scale"),
            layerTypes: tc.strings("layer_types"),
            numGlobalKeyValueHeads: json.int("num_global_key_value_heads") ?? tc.int("num_global_key_value_heads"),
            enableMoEBlock: tc.bool("enable_moe_block"),
            numExperts: tc.int("num_experts"),
            topKExperts: tc.int("top_k_experts"),
            moeIntermediateSize: tc.int("moe_intermediate_size"),
            attentionKEqualsV: json.bool("attention_k_eq_v") ?? tc.bool("attention_k_eq_v")
        )
    }
}

// ── JSON helpers ──────────────────────────────────────

extension Dictionary where Key == String, Value == Any {
    func string(_ key: String) -> String? { self[key] as? String }
    func int(_ key: String) -> Int? {
        if let v = self[key] as? Int { return v }
        if let v = self[key] as? Double { return Int(v) }
        if let v = self[key] as? NSNumber { return v.intValue }
        return nil
    }
    func float(_ key: String) -> Float? {
        if let v = self[key] as? Float { return v }
        if let v = self[key] as? Double { return Float(v) }
        if let v = self[key] as? NSNumber { return v.floatValue }
        return nil
    }
    func bool(_ key: String) -> Bool? {
        if let v = self[key] as? Bool { return v }
        return (self[key] as? NSNumber)?.boolValue
    }
    func strings(_ key: String) -> [String]? {
        self[key] as? [String]
    }
}