import ComposableArchitecture
import BackendEnvironmentClient
import Foundation
import GenAPI
import SunoModelClient
import Utilities

public struct Me: Codable, Equatable {
    public var models: [SunoModel]
    public var roles: [RoleKey: Bool]
    public var flags: [FlagKey: Bool]
    public var user: User
    public var statsigCustomProperties: StatsigCustomProperties?

    public init(models: [SunoModel], roles: [RoleKey: Bool], flags: [FlagKey: Bool], user: User, statsigCustomProperties: StatsigCustomProperties? = nil) {
        self.models = models
        self.roles = roles
        self.flags = flags
        self.user = user
        self.statsigCustomProperties = statsigCustomProperties
    }
}

public struct User: Codable, Equatable, Identifiable {
    enum CodingKeys: String, CodingKey {
        case id
        case handle
        case email
        case username
        case displayName = "display_name"
        case profileDescription = "profile_description"
        case avatarImageUrl = "avatar_image_url"
        case isHandleUpdated = "is_handle_updated"
        case phoneNumber = "phone_number"
        case inviteData = "ios_invite_data"
    }

    public var shareURL: URL {
        let baseUrl = BackendEnvironmentProvider.currentConfiguration().webEndpointHost
        return URL(string: "https://\(baseUrl)/@\(handle)")!
    }

    public var id: String
    public var handle: String
    public var email: String
    public var username: String
    public var displayName: String?
    public var profileDescription: String?
    public var avatarImageUrl: String?
    public var isHandleUpdated: Bool
    public var phoneNumber: String?
    public var inviteData: InviteData?

    /// For User.mock()
    init(id: String, handle: String, email: String, username: String, isHandleUpdated: Bool, inviteData _: InviteData? = nil) {
        self.id = id
        self.handle = handle
        self.email = email
        self.username = username
        self.isHandleUpdated = isHandleUpdated
    }

    /// From PATCH /api/profile
    public init?(_ genAPI: GetUserMeResponse, previousUser: User) {
        // TODO: make the userId non-optional from API. return username and inviteData as well
        guard let id = genAPI.userId else { return nil }
        self.id = id
        self.handle = genAPI.handle ?? ""
        self.email = genAPI.email ?? ""
        self.username = previousUser.username
        self.displayName = genAPI.displayName
        self.profileDescription = previousUser.profileDescription
        self.avatarImageUrl = genAPI.avatarImageUrl
        self.isHandleUpdated = genAPI.isHandleUpdated
        self.phoneNumber = genAPI.phoneNumber
        self.inviteData = previousUser.inviteData
    }
}

public struct InviteData: Codable, Equatable {
    enum CodingKeys: String, CodingKey {
        case canLogIn = "can_log_in_to_ios"
        case promoCodeUrl = "promo_code_url"
        case redemptionsLeft = "redemptions_left"
        case inviterHandle = "inviter_handle"
        case errorMsg = "error_msg"
    }

    public var canLogIn: Bool
    public var promoCodeUrl: String?
    public var redemptionsLeft: Int?
    public var inviterHandle: String?
    public var errorMsg: String?
}

public struct SunoModel: Codable, Equatable {
    enum CodingKeys: String, CodingKey {
        case id
        case name
        case majorVersion = "major_version"
        case externalKey = "external_key"
        case description
        case canUse = "can_use"
        case capabilities
        case maxLengths = "max_lengths"
    }

    public var id: String?
    public var name: String
    public var majorVersion: Int
    public var externalKey: String
    public var description: String

    // MARK: 4.5 Launch

    /// 4.5 launch model capabilities here are marked as `Optional` as they are flagged during rollout and may not be available.
    /// Make these non-optional in a future release.
    /// We should probably update the API to use `APIClientV2`...

    /// Can the user use this model (we still display it)
    /// https://github.com/suno-ai/glockenspiel/blob/11fd07dd533d279b02155da784d2a589ebf2db7e/studio_api/schema.json#L2548
    public var canUse: Bool?
    /// Model capabilities - determines what operations this model can perform
    /// https://github.com/suno-ai/glockenspiel/blob/11fd07dd533d279b02155da784d2a589ebf2db7e/studio_api/schema.json#L2553
    public var capabilities: [ModelCapability]?
    public var maxLengths: ModelTypeMaxLengths?
}

public extension SunoModel {
    var asSunoModelMetadata: SunoModelMetaData? {
        guard let id else {
            reportIssue("Found nil id from SunoModel \(self)")
            return nil
        }
        return .init(
            id: id,
            name: name,
            majorVersion: majorVersion,
            externalKey: externalKey,
            description: description,
            canUse: canUse,
            maxLengths: maxLengths?.asSunoModelMetadataMaxLengths
        )
    }
}

public extension SunoModelMetaData {
    init?(_ remote: GenAPI.ExternalModelTypeSchema) {
        // TODO: This is optional from the server, but we treat it as non-optional
        guard let id = remote.id?.uuidString else {
            reportIssue("Found nil id from SunoModel \(remote)")
            return nil
        }

        self.init(
            id: id,
            name: remote.name,
            majorVersion: remote.majorVersion,
            externalKey: remote.externalKey,
            description: remote.description,
            canUse: remote.canUse,
            maxLengths: .init(optionalRemote: remote.maxLengths),
            isDefaultModel: remote.isDefaultModel,
            isDefaultFreeModel: remote.isDefaultFreeModel,
            badges: remote.badges?.map {
                switch $0 {
                case .pro:
                    return .pro
                case .beta:
                    return .beta
                default:
                    return .other($0.rawValue)
                }
            },
            modelSelectorBadgeStyle: nil
        )
    }
}

public extension SunoModelMetaData.MaxLengths {
    init?(optionalRemote: GenAPI.ModelTypeMaxLengths?) {
        guard let remote = optionalRemote else { return nil }
        self.init(
            gptDescriptionPrompt: remote.gptDescriptionPrompt,
            negativeTags: remote.negativeTags,
            prompt: remote.prompt,
            tags: remote.tags,
            title: remote.title
        )
    }
}

/// Model capability for determining what operations a model can perform
public struct ModelCapability: RawRepresentable, Codable, Equatable, Hashable {
    public static let all = Self(rawValue: "all")
    public static let generate = Self(rawValue: "generate")
    public static let infill = Self(rawValue: "infill")
    public static let fixedInfill = Self(rawValue: "fixed_infill")
    public static let infillIntro = Self(rawValue: "infill_intro")
    public static let infillOutro = Self(rawValue: "infill_outro")
    public static let artistInfill = Self(rawValue: "artist_infill")
    public static let coverInfill = Self(rawValue: "cover_infill")
    public static let artistConsistency = Self(rawValue: "artist_consistency")
    public static let artistCover = Self(rawValue: "artist_cover")
    public static let artistExtend = Self(rawValue: "artist_extend")
    public static let cover = Self(rawValue: "cover")
    public static let upsample = Self(rawValue: "upsample")
    public static let extend = Self(rawValue: "extend")
    public static let uploadExtend = Self(rawValue: "upload_extend")
    public static let coverExtend = Self(rawValue: "cover_extend")
    public static let short = Self(rawValue: "short")
    public static let generateStem = Self(rawValue: "generate_stem")

    public var rawValue: String

    public init(rawValue: String) {
        self.rawValue = rawValue
    }
}

/// https://github.com/suno-ai/glockenspiel/blob/11fd07dd533d279b02155da784d2a589ebf2db7e/studio_api/schema.json#L5840
public struct ModelTypeMaxLengths: Codable, Equatable {
    enum CodingKeys: String, CodingKey {
        case gptDescriptionPrompt = "gpt_description_prompt"
        case negativeTags = "negative_tags"
        case prompt
        case tags
        case title
    }

    public let gptDescriptionPrompt: Int
    public let negativeTags: Int
    public let prompt: Int
    public let tags: Int
    public let title: Int
}

public struct FlagKey: RawRepresentable, Codable, Hashable, CodingKeyRepresentable {
    public static var support3ds: FlagKey { .init(rawValue: "support_3ds") }
    public static var newUpgradeMethod: FlagKey { .init(rawValue: "new_upgrade_method") }
    public static var continueAnywhere: FlagKey { .init(rawValue: "continue_anywhere") }
    public static var v3Alpha: FlagKey { .init(rawValue: "v3_alpha") }
    public static var playlists: FlagKey { .init(rawValue: "playlists") }
    public static var profilesEdit: FlagKey { .init(rawValue: "profiles-edit") }
    public static var homeExplorePlaylists: FlagKey { .init(rawValue: "home-explore-playlists") }
    public static var iosSubscriptions: FlagKey { .init(rawValue: "ios-subscriptions") }
    public static var iosV1Navigation: FlagKey { .init(rawValue: "ios-v1-navigation") }
    public static var iosV2ClipCreation: FlagKey { .init(rawValue: "ios-v2-clip-creation") }
    public var rawValue: String

    public init(rawValue: String) {
        self.rawValue = rawValue
    }
}

public struct RoleKey: RawRepresentable, Codable, Hashable, CodingKeyRepresentable {
    public static var hasAcceptedCustomModeTOS: RoleKey { .init(rawValue: "has_accepted_custom_mode_tos") }
    public static var sub: RoleKey { .init(rawValue: "sub") }
    public static var isDayZeroUser: RoleKey { .init(rawValue: "is_day_zero_user") }

    public var rawValue: String

    public init(rawValue: String) {
        self.rawValue = rawValue
    }
}

public struct StatsigCustomProperties: Codable, Equatable {
    public let custom: [String: String]?
    public let customIDs: [String: String]?
}

public extension StatsigCustomProperties {
    init?(_ remote: GenAPI.StatsigCustomProperties?) {
        guard let remote else { return nil }
        self.custom = remote.custom
        self.customIDs = remote.customIds
    }
}

public extension User {
    static func mock(
        id: String = UUID().uuidString
    ) -> User {
        User(id: id, handle: "app-ios", email: "ios@suno.com", username: "app-ios", isHandleUpdated: true)
    }

    static func empty() -> User {
        User(id: UUID().uuidString, handle: "", email: "ios@suno.com", username: "", isHandleUpdated: false)
    }
}

extension Me: Then {}

// MARK: APIClientV2

public extension ModelTypeMaxLengths {
    init(_ remote: GenAPI.ModelTypeMaxLengths) {
        self.gptDescriptionPrompt = remote.gptDescriptionPrompt
        self.negativeTags = remote.negativeTags
        self.prompt = remote.prompt
        self.tags = remote.tags
        self.title = remote.title
    }

    var asSunoModelMetadataMaxLengths: SunoModelMetaData.MaxLengths {
        .init(
            gptDescriptionPrompt: gptDescriptionPrompt,
            negativeTags: negativeTags,
            prompt: prompt,
            tags: tags,
            title: title
        )
    }
}

public extension SunoModel {
    init(_ remote: GenAPI.ExternalModelTypeSchema) {
        // TODO: This is optional from the server, but we treat it as non-optional
        self.id = remote.id?.uuidString
        self.name = remote.name
        self.majorVersion = remote.majorVersion
        // TODO: Also handle this optional case :/
        self.externalKey = remote.externalKey
        self.description = remote.description
        self.canUse = remote.canUse
        self.capabilities = remote.capabilities?.map { ModelCapability(rawValue: $0.rawValue) }
        if let maxLengths = remote.maxLengths {
            self.maxLengths = .init(maxLengths)
        }
    }
}

public extension User {
    init(_ remote: GenAPI.LoggedInSessionUser) throws {
        // TODO: make the userId non-optional from API.
        guard let id = remote.id?.uuidString else {
            throw DecodingError.typeMismatch(User.self, .init(codingPath: [], debugDescription: "Missing required uuid while decoding User from \(remote)"))
        }
        self.id = id.lowercased()
        self.handle = remote.handle ?? ""
        self.email = remote.email
        self.username = remote.username
        self.displayName = remote.displayName
        self.profileDescription = remote.profileDescription
        self.avatarImageUrl = remote.avatarImageUrl
        self.isHandleUpdated = remote.isHandleUpdated ?? false
        self.phoneNumber = remote.phoneNumber
    }
}

public extension Me {
    init(_ remote: GenAPI.LoggedInSessionResponse) throws {
        self.user = try .init(remote.user)
        self.models = remote.models.map { .init($0) }
        self.roles = remote.roles.reduce(into: [:]) { result, role in
            result[RoleKey(rawValue: role.key)] = role.value.boolValue
        }
        self.flags = remote.flags.reduce(into: [:]) { result, flag in
            result[FlagKey(rawValue: flag.key)] = flag.value.boolValue
        }
        self.statsigCustomProperties = .init(remote.statsigCustomProperties)
    }

    init(_ remote: GenAPI.UnauthenticatedSessionResponse) throws {
        self.user = try remote.user.to(User.self)
        self.models = remote.models.map { .init($0) }
        self.roles = remote.roles.reduce(into: [:]) { result, role in
            result[RoleKey(rawValue: role.key)] = role.value.boolValue
        }
        self.flags = remote.flags.reduce(into: [:]) { result, flag in
            result[FlagKey(rawValue: flag.key)] = flag.value.boolValue
        }
        self.statsigCustomProperties = .init(remote.statsigCustomProperties)
    }
}
