import ComposableArchitecture
import Foundation
import StatsigClient
import UserDefaultsClient
import Utilities

let log = Logger(category: "SunoModelClient")

@DependencyClient
public struct SunoModelClient {
    public var configure: (_ sunoModels: [SunoModelMetaData], _ userAccess: SunoModelUserAccess, _ userID: String, _ remasterModels: [RemasterModelMetaData]?) async -> Void

    public var getCurrentModelExternalKey: (_ isAudioUpload: Bool) -> String = { _ in "chirp-v3-5" }
    public var getCurrentModelName: () -> String?
    public var getCurrentModelMarketingLevelUnderstanding: () -> SunoModelMetaData.MarketingLevelUnderstanding = { .previousToV4 }

    // Remaster model methods
    public var getCurrentRemasterModelExternalKey: () -> String = { ModelConstants.defaultModel.rawValue }
    public var getAvailableRemasterModels: () -> [RemasterModelMetaData] = { [] }
    public var setRemasterModel: (_ externalKey: String) async -> Void

    public var setUserAccess: (_ userAccess: SunoModelUserAccess) -> Void
    public var toggleModelOverride: () async -> Void
    public var clearModelOverride: () async -> Void

    public var setModel: (_ id: String) async -> Void
    public var setModelWithMarketingLevelUnderstanding: (_ marketingLevelUnderstanding: SunoModelMetaData.MarketingLevelUnderstanding) async -> Void
}

extension SunoModelClient: DependencyKey {
    public static let liveValue: Self = {
        @Dependency(UserDefaultsClient.self) var userDefaults

        /** Only ever write to this from Suno Model Client */
        @Shared(.inMemory(.selectedSunoModel)) var selectedSunoModel: SunoModelMetaData = .modelDefault
        @Shared(.inMemory(.selectedSunoRemasterModel)) var selectedSunoRemasterModel: RemasterModelMetaData?
        @Shared(.inMemory(.sunoModelUserAccessCategory)) var sunoModelUserAccessCategory: SunoModelUserAccess = .defaultLimitedAccess
        @Shared(.inMemory(.availableModels)) var availableModels: [SunoModelMetaData] = []

        var cachedUserID: String = ""
        var currentUserModels: [SunoModelMetaData] = []
        var currentRemasterModels: [RemasterModelMetaData] = []

        func internal_recomputeAvailableModels() -> [SunoModelMetaData] {
            currentUserModels.map {
                /// Free gens are our last "hacky" part of this pipeline.
                /// We may not be supporting free gens long-term, so this is staying as-is for now. The correct way to support free gens would be to pull this value from the backend per-model as well.
                if $0.majorVersion == 4 {
                    var model = $0
                    if case let .freeUserWithFreeLifetimeGenerationsLeft(freeV4Gens) = sunoModelUserAccessCategory.v4Access {
                        model.songsLeft = freeV4Gens
                    }
                    return model
                } else if $0.name == "v4.5+" || $0.externalKey == "chirp-bluejay", FeatureFlag.legacy.bluejayModelHighlighted {
                    // Fragile, but this is so that we can highlight this model in pink for the v4.5+ launch
                    // Unfortunately, the major version of this model is the same as the major version of the v4.5 model, so we have to match on the model string.
                    // Hooked up to a Statsig so that we can turn this off from the dashboard.
                    var model = $0
                    if !FeatureFlag.legacy.v5Launch {
                        model.isNew = true
                    }
                    return model
                } else if $0.name == "v5" || $0.externalKey == "chirp-crow", FeatureFlag.legacy.v5Launch {
                    // Fragile, but this is so that we can highlight this model in pink for the v5 launch
                    // Unfortunately, the major version of this model is the same as the major version of the v4.5 model, so we have to match on the model string.
                    // Hooked up to a Statsig so that we can turn this off from the dashboard.
                    var model = $0
                    model.isNew = true
                    return model
                } else {
                    return $0
                }
                // adding reversed here so that the consumers, CreatClip and CreateCover can show the models in ascending order
            }.reversed()
        }

        func internal_getCurrentModelOverride() -> SunoModelOverrideData {
            let currentOverridesData = userDefaults.dataForKey(.modelVersioningOverride)
            return SunoModelOverrideData.fromData(currentOverridesData)
        }

        func internal_getCurrentModel() -> SunoModelMetaData {
            let override = internal_getCurrentModelOverride()
            let model: SunoModelMetaData

            if let overrideModel = override.overrideMap[cachedUserID],
               // Make sure that we match against the backend provided models, otherwise go to defaults
               let overrideModelFromBackend = currentUserModels.first(where: { $0.id.lowercased() == overrideModel.id.lowercased() && $0.canUse != false })
            {
                model = overrideModelFromBackend
            } else {
                let modelData = userDefaults.dataForKey(.defaultSunoModelForUserAccessCategory)
                model = SunoModelMetaData.fromDataOrDefault(modelData)
            }

            return model
        }

        func internal_updateAppModelSourceOfTruth() {
            $selectedSunoModel.withLock { model in
                model = internal_getCurrentModel()
            }
        }

        func internal_getCurrentRemasterModelOverride() -> RemasterModelOverrideData {
            let currentOverridesData = userDefaults.dataForKey(.remasterModelVersioningOverride)
            return RemasterModelOverrideData.fromData(currentOverridesData)
        }

        func internal_getCurrentRemasterModel() -> RemasterModelMetaData? {
            let override = internal_getCurrentRemasterModelOverride()
            let model: RemasterModelMetaData?

            if let overrideModel = override.overrideMap[cachedUserID],
               // Make sure that we match against the backend provided models
               let overrideModelFromBackend = currentRemasterModels.first(where: { $0.externalKey == overrideModel.externalKey })
            {
                model = overrideModelFromBackend
            } else {
                let modelData = userDefaults.dataForKey(.defaultSunoRemasterModelForUserAccessCategory)
                model = RemasterModelMetaData.fromData(modelData)
            }

            return model
        }

        func internal_updateAppRemasterModelSourceOfTruth() {
            $selectedSunoRemasterModel.withLock { model in
                model = internal_getCurrentRemasterModel()
            }
        }

        func internal_updateCurrentUserAccess(_ userAccess: SunoModelUserAccess) {
            $sunoModelUserAccessCategory.withLock { category in
                category = userAccess
            }
        }

        func internal_isModelOverridden() -> Bool {
            var override = internal_getCurrentModelOverride()
            if let _ = override.overrideMap[cachedUserID] {
                return true
            } else {
                return false
            }
        }

        func internal_setModelOverride(_ modelOverride: SunoModelMetaData.MarketingLevelUnderstanding) async {
            let modelMetaData: SunoModelMetaData
            switch modelOverride {
            case .previousToV4, .v3Dot5:
                modelMetaData = .modelV3_5

            case .v4:
                modelMetaData = .modelV4

            case .auk:
                modelMetaData = currentUserModels.first(where: { $0.marketingLevelUnderstanding == .auk }) ?? .modelV4

            case .bluejay:
                modelMetaData = currentUserModels
                    .first(where: { $0.marketingLevelUnderstanding == .bluejay }) ?? .modelV4

            case .v5:
                modelMetaData = currentUserModels
                    .first(where: { $0.marketingLevelUnderstanding == .v5 }) ?? .modelV4
            }

            var override = internal_getCurrentModelOverride()
            override.overrideMap[cachedUserID] = modelMetaData

            await userDefaults.setData(override.asJSONData, .modelVersioningOverride)

            Task { @MainActor in
                internal_updateAppModelSourceOfTruth()
            }
        }

        func internal_setModel(_ id: String) async {
            guard let model = currentUserModels.first(where: { $0.id.lowercased() == id.lowercased() }) else {
                return
            }

            let modelMetaData: SunoModelMetaData = currentUserModels.first(where: { $0.id.lowercased() == id.lowercased() }) ?? .modelV4

            var override = internal_getCurrentModelOverride()
            override.overrideMap[cachedUserID] = modelMetaData

            await userDefaults.setData(override.asJSONData, .modelVersioningOverride)

            Task { @MainActor in
                internal_updateAppModelSourceOfTruth()
            }
        }

        func internal_setModelWithMarketingLevelUnderstanding(_ marketingLevelUnderstanding: SunoModelMetaData.MarketingLevelUnderstanding) async {
            guard let modelMetadata = currentUserModels.first(where: { $0.marketingLevelUnderstanding == marketingLevelUnderstanding }) else {
                // Log non-fatal
                log.telemetry.assertionFailure("Failed to set model using marketing level understanding: \(marketingLevelUnderstanding.rawValue). Available models are: \(currentUserModels)")
                return
            }

            var override = internal_getCurrentModelOverride()
            override.overrideMap[cachedUserID] = modelMetadata

            await userDefaults.setData(override.asJSONData, .modelVersioningOverride)

            Task { @MainActor in
                internal_updateAppModelSourceOfTruth()
            }
        }

        func internal_clearModelOverride() async {
            var override = internal_getCurrentModelOverride()
            override.overrideMap.removeValue(forKey: cachedUserID)
            await userDefaults.setData(override.asJSONData, .modelVersioningOverride)

            Task { @MainActor in
                internal_updateAppModelSourceOfTruth()
            }
        }

        func internal_setDefaultModel() async {
            /**
                This is the default model that is used for this tier.
                This does not need to check overrides.
                Generate this using the *highest* model with `is_default_model == true`. For added safety, we check `can_use == true` or `can_use == nil` model available to the user (We have to treat `nil` as `true` right now...), but this behavior is considered to be undefined by the backend.
             */
            // Find the highest version that the user can use
            let highestAvailableDefaultVersion = currentUserModels
                .filter { $0.isDefaultModel == true && $0.canUse != false }
                .max { $0.majorVersion < $1.majorVersion } ?? SunoModelMetaData.modelV3_5 // Default to 3.5 as a fallback
            await userDefaults.setData(highestAvailableDefaultVersion.asJSONData, .defaultSunoModelForUserAccessCategory)
        }

        func internal_setDefaultRemasterModel() async {
            // Find the default remaster model or use the first one available
            let defaultRemasterModel = currentRemasterModels.first(where: { $0.isDefaultModel }) ?? currentRemasterModels.first
            await userDefaults.setData(defaultRemasterModel?.asJSONData, .defaultSunoRemasterModelForUserAccessCategory)
        }

        func internal_setRemasterModel(_ externalKey: String) async {
            guard let model = currentRemasterModels.first(where: { $0.externalKey == externalKey }) else {
                return
            }

            var override = internal_getCurrentRemasterModelOverride()
            override.overrideMap[cachedUserID] = model

            await userDefaults.setData(override.asJSONData, .remasterModelVersioningOverride)

            Task { @MainActor in
                internal_updateAppRemasterModelSourceOfTruth()
            }
        }

        func internal_clearRemasterModelOverride() async {
            var override = internal_getCurrentRemasterModelOverride()
            override.overrideMap.removeValue(forKey: cachedUserID)
            await userDefaults.setData(override.asJSONData, .remasterModelVersioningOverride)

            Task { @MainActor in
                internal_updateAppRemasterModelSourceOfTruth()
            }
        }

        return Self(
            configure: { userSunoModels, userAccess, userID, remasterModels in
                currentUserModels = userSunoModels
                cachedUserID = userID

                // Store remaster models and set default if available
                if let remasterModels = remasterModels, !remasterModels.isEmpty {
                    currentRemasterModels = remasterModels
                    await internal_setDefaultRemasterModel()
                }

                await internal_setDefaultModel()

                // Compute and store available models for shared access
                $availableModels.withLock { models in
                    models = internal_recomputeAvailableModels()
                }

                Task { @MainActor in
                    internal_updateCurrentUserAccess(userAccess)
                    internal_updateAppModelSourceOfTruth()
                    internal_updateAppRemasterModelSourceOfTruth()
                }
            },
            getCurrentModelExternalKey: { isAudioUpload in
                /** Prefer a single return statement here */
                let model = internal_getCurrentModel()
                return model.getExternalKey(isAudioUpload)
            },
            getCurrentModelName: {
                let model = internal_getCurrentModel()
                return model.name
            },
            getCurrentModelMarketingLevelUnderstanding: {
                let model = internal_getCurrentModel()
                return model.marketingLevelUnderstanding
            },
            getCurrentRemasterModelExternalKey: {
                // Return current remaster model or fallback to default
                let model = internal_getCurrentRemasterModel()
                return model?.externalKey ?? ModelConstants.defaultModel.rawValue
            },
            getAvailableRemasterModels: {
                currentRemasterModels
            },
            setRemasterModel: { externalKey in
                await internal_setRemasterModel(externalKey)
            },
            setUserAccess: { updatedUserAccess in
                // Capture in order to work across actor boundaries
                let currentSunoModelUserAccess = sunoModelUserAccessCategory
                Task { @MainActor in
                    // if the user just ran out of free v4 generations, update the default model
                    if updatedUserAccess.v4Access != currentSunoModelUserAccess.v4Access,
                       case .freeUserWithoutFreeLifetimeGenerationsLeft = updatedUserAccess.v4Access
                    {
                        currentUserModels = currentUserModels.map {
                            if $0.majorVersion == 4 {
                                var model = $0
                                model.canUse = false
                                return model
                            } else {
                                return $0
                            }
                        }
                        await internal_setDefaultModel()
                        internal_updateAppModelSourceOfTruth()
                        // Update available models to reflect the new canUse status
                        $availableModels.withLock { models in
                            models = internal_recomputeAvailableModels()
                        }
                    }
                    internal_updateCurrentUserAccess(updatedUserAccess)
                }
            },
            toggleModelOverride: {
                if internal_isModelOverridden() {
                    await internal_clearModelOverride()
                } else {
                    await internal_setModelOverride(.v3Dot5)
                }
            },
            clearModelOverride: {
                await internal_clearModelOverride()
            },
            setModel: { modelId in
                await internal_setModel(modelId)
            },
            setModelWithMarketingLevelUnderstanding: { marketingLevelUnderstanding in
                await internal_setModelWithMarketingLevelUnderstanding(marketingLevelUnderstanding)
            }
        )
    }()
}

extension SunoModelClient: TestDependencyKey {
    public static let previewValue = Self()

    public static let testValue = Self()
}

public extension SunoModelClient {
    enum ModelConstants: String {
        case defaultModel = "chirp-bass"
    }
}