import BackendEnvironmentClient
import Combine
import ComposableArchitecture
import Foundation
import InstanceIDClient
import Statsig
import Utilities

@DependencyClient
public struct StatsigClient {
    public var startStatsig: (@escaping (Error?) -> Void) -> Void

    /// Call this once the user has logged in
    /// There's no way to log a user out with Statsig
    public var setUser: (User) async throws -> Void
    public var getStableID: () -> String?

    public var trackEvent: (StatigEvent) -> Void

    public var state: @Sendable () -> AsyncStream<StatsigClient.State> = { .never }
    public var currentState: () -> StatsigClient.State = { .uninitialized }

    /// Use `FeatureFlag.{ParameterStoreName}.{key}` instead of calling this directly
    var getValue: (_ parameterStore: String, _ key: String, _ defaultValue: StatsigDynamicConfigValue) -> StatsigDynamicConfigValue?
    /// Use `FeatureFlag.{ParameterStoreName}.markExposed(\.${key})` instead of calling this directly
    var markExposure: (_ parameterStore: String, _ key: String, _ defaultValue: StatsigDynamicConfigValue) -> Void

    /// Get dynamic config values
    public var getDynamicConfig: (_ configName: String, _ defaultValue: StatsigDynamicConfigValue) -> StatsigDynamicConfigValue?

    /// Clear the UserID and user properties we store in `@Shared(.appStorage(.lastStatsigUserId))` and related keys when signing out
    public var clearCachedUser: () -> Void
}

public extension StatsigClient {
    enum State {
        case uninitialized
        case initializing
        case initialized
        case initializedWithUser
    }

    enum User: Equatable {
        case authenticated(
            userId: String,
            email: String?,
            custom: [String: String]?,
            customIDs: [String: String]?
        )
        case anonymous

        public var userId: String? {
            switch self {
            case .authenticated(userId: let id, email: _, custom: _, customIDs: _): id
            case .anonymous: nil
            }
        }
    }
}

public enum StatigEvent {
    case onboardingStarted(stableID: String?, sunoUserID: String)

    var eventName: String {
        let mirror = Mirror(reflecting: self)
        return mirror.children.isEmpty ? "\(self)" : mirror.children.first!.label ?? String(describing: mirror.displayStyle)
    }

    var metadata: [String: String]? {
        switch self {
        case let .onboardingStarted(stableID, sunoUserID):
            [
                "stable_id": stableID ?? "",
                "suno_user_id": sunoUserID,
            ]
        }
    }
}

protocol StatsigParameterStore {
    /// Needs to exactly match the name in the dashboard
    static var parameterStoreName: String { get }
}

struct StatsigParameterValue<T: StatsigDynamicConfigValue>: StatsigParameterValueType {
    typealias RawValue = String

    let name: String
    let defaultLocalValue: T
}

protocol StatsigParameterValueType {
    associatedtype ParameterValue: StatsigDynamicConfigValue

    var name: String { get }
    var defaultLocalValue: ParameterValue { get }
}

extension StatsigClient: DependencyKey {
    /// https://console.statsig.com/64RBMXCoSmsTc9oTU9ghAk/param_stores/suno-ios
    private static let parameterStoreName = "suno-ios"

    public static let liveValue: Self = {
        @Dependency(BackendEnvironmentClient.self) var backendEnvironment
        @Dependency(\.telemetryClient) var telemetry
        @Dependency(\.instanceIdClient) var instanceIdClient
        @Shared(.appStorage(.lastStatsigUserId)) var cachedStatsigUserId: String = ""
        @Shared(.appStorage(.lastStatsigUserUSState)) var cachedStatsigUserUSState: String = ""

        let environmentConfiguration = backendEnvironment.configuration()

        guard let info = Bundle.main.infoDictionary,
              let clientKey = info["STATSIG_CLIENT_KEY"] as? String,
              !clientKey.isEmpty else {
            fatalError("STATSIG_CLIENT_KEY is not defined in Info.plist")
        }

        let statsigEnvironment = switch environmentConfiguration.statsigEnvironment {
        case "staging":
            StatsigEnvironment(tier: .Staging)
        case "production":
            StatsigEnvironment(tier: .Production)
        default:
            StatsigEnvironment(tier: .Development)
        }

        /// a timeout of 0 tells Statsig to keep trying to fetch from the network until it succeeds
        let statsigOptions = StatsigOptions(
            initTimeout: 0,
            environment: statsigEnvironment,
            overrideStableID: instanceIdClient.instanceId()
        )

        let anonymousUserId = instanceIdClient.instanceId()

        let initialState = StatsigClient.State.uninitialized
        let stateStorage = LockIsolated<StatsigClient.State>(initialState)
        let (stateStream, stateContinuation) = AsyncStream.makeStream(of: StatsigClient.State.self)
        stateContinuation.yield(initialState)

        func setCachedStatsigUserId(_ userId: String?) {
            $cachedStatsigUserId.withLock { $0 = userId ?? "" }
        }

        func setCachedStatsigUserUSState(_ usState: String?) {
            $cachedStatsigUserUSState.withLock { $0 = usState ?? "" }
        }

        func updateStatsigUser(
            user: StatsigUser,
            completion: @escaping (Error?) -> Void
        ) {
            /// Statsig doesn't throw, but it's clearer if we do and choose to handle it well
            /// This async also turns the Statsig updateUser event into a potentially blocking call,
            /// but sometimes we need to make a check right after updating user
            Statsig.updateUserWithResult(user) { optionalError in
                if let error = optionalError {
                    telemetry.record(error: error, message: "Couldn't update user")
                    completion(error)
                } else {
                    telemetry.record(message: "Statsig: user updated successfully. user:\(user)")

                    let state = StatsigClient.State.initializedWithUser
                    stateStorage.withValue { $0 = state }
                    stateContinuation.yield(state)

                    persistCriticalFlags()
                    completion(nil)
                }
            }
        }

        func persistCriticalFlags() {
            /*
             We need to persist this flag because we need access to it before Statsig
             is initialized. Statsig persists these under the hood, but we only get the
             default value (false, in this case) if we try accessing it before the SDK
             is initialized - bit of a chicken vs egg problem.
             */
            @Shared(.appStorage(.startUpOrchestratorEnabled)) var startUpOrchestratorEnabled: Bool = false
            $startUpOrchestratorEnabled.withLock { $0 = FeatureFlag.general.startUpOrchestratorEnabled }
        }

        return Self(
            startStatsig: { callback in
                var hasEmittedInitializing = false

                let cachedUserId = cachedStatsigUserId
                // TODO: (BA) Cache the User instead of different components separately.
                // Ideally, this is part of a slighter larger refactor to make this flow more
                // robust for app launch experiments.
                let cachedUSState = cachedStatsigUserUSState
                let initialUser: StatsigUser? = {
                    guard !cachedUserId.isEmpty else { return nil }
                    if !cachedUSState.isEmpty {
                        return StatsigUser(
                            userID: cachedUserId,
                            email: nil,
                            custom: ["us_state": cachedUSState],
                            customIDs: ["US State": cachedUSState]
                        )
                    } else {
                        return StatsigUser(userID: cachedUserId)
                    }
                }()

                Statsig.initialize(
                    sdkKey: clientKey,
                    user: initialUser,
                    options: statsigOptions,
                    completion: { error in
                        // If we don't start with a cacched source, emit the initializing state now
                        if !hasEmittedInitializing {
                            stateStorage.withValue { $0 = .initializing }
                            stateContinuation.yield(.initializing)
                        }

                        stateStorage.withValue { $0 = .initialized }
                        stateContinuation.yield(.initialized)
                        callback(error)
                    }
                )

                let source = Statsig.getInitializeResponseJson().evaluationDetails.source
                switch source {
                case .Cache:
                    // Emit initializing state after calling initialize, if we do have cached values ready
                    // Calling Statsig.initialize sets this source to cache on all but the first app open
                    stateStorage.withValue { $0 = .initializing }
                    stateContinuation.yield(.initializing)
                    hasEmittedInitializing = true

                default:
                    break
                }
            },
            setUser: { user in
                let statsigUser = user.toStatsigUser(
                    anonymousUserId: anonymousUserId,
                    locale: .current
                )

                try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
                    switch user {
                    case .anonymous:
                        telemetry.record(message: "Statsig: setting anonymous user. userID:\(anonymousUserId)")
                        setCachedStatsigUserId(nil)
                        setCachedStatsigUserUSState(nil)

                    case let .authenticated(userId, email, custom, customIDs):
                        telemetry
                            .record(
                                message: "Statsig: setting user. userID:\(userId), email:\(email ?? ""), custom: \(String(describing: custom)), customIDs: \(String(describing: customIDs))"
                            )

                        setCachedStatsigUserId(userId)
                        let usState = custom?["us_state"] ?? customIDs?["US State"]
                        setCachedStatsigUserUSState(usState)
                    }

                    updateStatsigUser(user: statsigUser) { error in
                        if let error {
                            continuation.resume(throwing: error)
                        } else {
                            continuation.resume()
                        }
                    }
                }
            },
            getStableID: {
                Statsig.getStableID()
            },
            trackEvent: { event in
                Statsig.logEvent(event.eventName, metadata: event.metadata)
            },
            state: {
                stateStream
            },
            currentState: {
                stateStorage.value
            },
            getValue: { parameterStoreName, key, defaultValue in
                let paramStore = Statsig.getParameterStoreWithExposureLoggingDisabled(parameterStoreName)
                /// Though the docs for this function say "If a valid value is found, a layer exposure event will be fired", an exposure event is not fired since we used `getParameterStoreWithExposureLoggingDisabled` to create this store
                /// Docstrings are wrong from Statsig for this
                return paramStore.getValue(forKey: key, defaultValue: defaultValue)
            },
            markExposure: { parameterStore, key, defaultValue in
                let paramStore = Statsig.getParameterStore(parameterStore)
                /// Unfortunately, there's no easy way to directly log an exposure for a `parameterStore` value.
                /// There are 4 ways to log an exposure `manuallyLog{Gate,Config,Experiment,LayerParameter}Exposure`
                /// Since a parameter store value could be backed by any one of those 4 values or a static value, the simplest solution is to get a `getParameterStore` that will mark an exposure as a side effect of getting a value
                let result = paramStore.getValue(forKey: key, defaultValue: defaultValue)
                telemetry.record(message: "Marking exposure for parameterStore=\(parameterStore) key=\(key) value=\(String(describing: result))")
            },
            getDynamicConfig: { configName, _ in
                let config = Statsig.getConfig(configName)
                return config.value
            },
            clearCachedUser: {
                setCachedStatsigUserId(nil)
                setCachedStatsigUserUSState(nil)
            }
        )
    }()
}

extension StatsigClient: TestDependencyKey {
    public static let previewValue: Self = Self.noop

    public static let testValue: Self = Self(
        startStatsig: { callback in
            callback(nil)
        },
        setUser: { _ in },
        getStableID: { .none },
        trackEvent: { _ in },
        state: { .never },
        currentState: { .uninitialized },
        getValue: { _, _, _ in false },
        markExposure: { _, _, _ in },
        getDynamicConfig: { _, defaultValue in defaultValue },
        clearCachedUser: { }
    )
}

public extension StatsigClient {
    static let noop = StatsigClient(
        startStatsig: { _ in },
        setUser: { _ in },
        getStableID: { .none },
        trackEvent: { _ in },
        state: { .never },
        currentState: { .uninitialized },
        getValue: { _, _, _ in false },
        markExposure: { _, _, _ in },
        getDynamicConfig: { _, defaultValue in defaultValue },
        clearCachedUser: { }
    )
}

// MARK: - Domain User -> Statsig User

private extension StatsigClient.User {
    func toStatsigUser(
        anonymousUserId: String,
        locale: Locale
    ) -> StatsigUser {
        let statsigCountry = locale.region?.identifier

        switch self {
        case .anonymous:
            return StatsigUser(
                userID: anonymousUserId,
                country: statsigCountry
            )

        case .authenticated(userId: let userId, email: let email, custom: let custom, customIDs: let customIDs):
            return StatsigUser(
                userID: userId,
                email: email,
                country: statsigCountry,
                custom: custom,
                customIDs: customIDs
            )
        }
    }
}