import APIClient
import Combine
import ComposableArchitecture
import Foundation
import HCaptchaClient
import Localization
import SunoModelClient
import Utilities

@DependencyClient
public struct BlendedCreateClient {
    // MARK: - API

    // Create hooks
    // Return a `BlendedCreatePrompt` because this function transforms the prompt before creation
    public var generate: (BlendedCreatePrompt, _ createType: CreateType, _ lyricsModel: String?) -> BlendedCreatePrompt = { _, _, _ in .init() }
    public var uploadAudio: @Sendable (AudioRecording) -> AsyncThrowingStream<AudioUploadStatus, Error> = { _ in .never }
    public var initializeClip: @Sendable (_ uploadRequestId: String) async throws -> String

    // Generate more - direct async method that doesn't use the shared stream
    public var generateMore: @Sendable (BlendedCreatePrompt, _ createType: CreateType, _ lyricsModel: String?) async throws -> [Clip] = { _, _, _ in [] }

    // Event Bus
    public var generateResultStream: () -> AsyncStream<BlendedGenerationResult> = { .never }

    // MARK: - Models

    // Generate
    public typealias BlendedGenerationResult = Result<[Clip], Error>
    // This guarantees at generation time that each of these create vectors has everything that they need
    public enum CreateType {
        case textOnly
        case audio(clipId: String)
    }

    // Audio Upload
    public enum AudioUploadStatus {
        case uploadEvent(APIClient.UploadEvent)
        case status(UploadRequestStatus)
    }

    public enum AudioUploadError: Error {
        case fileCreationFailed
        case verificationFailed
        case statusCheckFailed
    }

    public enum InitializeClipError: Error {
        case noClipIdReturned
    }
}

extension BlendedCreateClient: DependencyKey {
    public static let liveValue: Self = {
        @Dependency(SunoModelClient.self) var sunoModelClient
        @Dependency(APIClientV2.self) var apiClientV2
        @Dependency(APIClient.self) var apiClient
        @Dependency(\.continuousClock) var clock
        @Dependency(HCaptchaClient.self) var hCaptchaClient

        struct AudioUploadCancellable: Hashable {}

        // Should move this into a "Clip Generation" client
        let subject = PassthroughSubject<BlendedGenerationResult, Never>()

        // Helpers
        func _generateWithHCaptchaRetries(_ prompt: BlendedCreatePrompt, model: String, lyricsModel: String?) async throws -> [Clip] {
            do {
                // Generate
                return try await _generate(prompt, model: model, lyricsModel: lyricsModel)
            } catch {
                log.telemetry.error(error)
                // Handle HCaptcha retries
                // Simple logic for now just to retry once
                switch error {
                case let error as APIError:
                    if case APIError.invalidHCaptchaToken = error {
                        log.info("Invalid HCaptcha token, retrying the token fetch and generation in-client.")
                        // Only try once for now
                        let token = try await hCaptchaClient.getToken()
                        var newPrompt = prompt
                        newPrompt.token = token
                        // Generate again
                        return try await _generate(newPrompt, model: model, lyricsModel: lyricsModel)
                    } else {
                        throw error
                    }

                default:
                    // Rethrow
                    throw error
                }
            }
        }

        func _generate(_ prompt: BlendedCreatePrompt, model: String, lyricsModel: String?) async throws -> [Clip] {
            switch prompt.createMode {
            case .simple:
                // Simple create - No lyrics or tricks
                if prompt.lyricsMode != .write {
                    let response = try await APIClientV2.underlying.send(
                        Paths.generate.v2.post(
                            prompt.asSimpleGenAPI(model: model, lyricsModel: lyricsModel)
                        )
                    )
                    return try response.value.clips.map(Clip.init)
                } else {
                    // Complex "Simple" create
                    let generateResponse = try await APIClientV2.underlying.send(
                        Paths.generate.v2.post(
                            prompt.asBlendedGenAPI(model: model, lyricsModel: lyricsModel)
                        )
                    )

                    return try generateResponse.value.clips.map(Clip.init)
                }

            case .custom:
                @Shared(.inMemory(.billingInfo)) var billingInfo: SubscriptionInfoResponse?
                // Custom create - Custom lyrics
                let response = try await APIClientV2.underlying.send(
                    Paths.generate.v2.post(
                        prompt.asCustomGenAPI(model: model, lyricsModel: lyricsModel, useSliders: billingInfo?.plan != nil)
                    )
                )
                return try response.value.clips.map(Clip.init)
            }
        }

        return Self(
            generate: { prompt, createType, lyricsModel in
                // We do some last minute write to the prompt before generation
                var mutablePrompt = prompt
                mutablePrompt.lyricsModel = lyricsModel ?? prompt.lyricsModel
                // Used for the model selector
                var isAudioUpload: Bool = false
                // Audio create mode transforms
                if case .audio(let clipId) = createType {
                    isAudioUpload = true
                    // Set the clipId to the proper fields
                    // This is done on-demand before generation because we don't want to
                    switch mutablePrompt.audioCreateStyle {
                    case .cover:
                        mutablePrompt.coverClipId = mutablePrompt.clipId
                        // Clean up the other task
                        mutablePrompt.continueClipId = nil
                        // Parity w/ web, it seems
                        mutablePrompt.continueAt = 0
                        mutablePrompt.task = .cover

                    case .extend:
                        mutablePrompt.continueClipId = mutablePrompt.clipId
                        // Clean up the other task
                        mutablePrompt.coverClipId = nil
                        // Huh?? This is from legacy but I don't know how this applies to when audioRecording is nil...
                        mutablePrompt.continueAt = mutablePrompt.audioRecording?.duration
                        // ...
                        mutablePrompt.task = .uploadExtend
                    }
                }

                // Fetch the model
                let model = sunoModelClient.getCurrentModelExternalKey(isAudioUpload: isAudioUpload)

                // For snappiness
                Task {
                    let result = await BlendedGenerationResult(catching: { try await _generateWithHCaptchaRetries(mutablePrompt, model: model, lyricsModel: lyricsModel) })
                    // Save clips to Saved Prompts
                    if case let .success(clips) = result {
                        @Shared(.fileStorage(.savedPrompts)) var savedPrompts: [Clip.ID: Prompt] = [:]
                        for clip in clips {
                            $savedPrompts.withLock {
                                $0[clip.id] = mutablePrompt.toPrompt()
                            }
                        }
                    }
                    subject.send(result)
                }
                // Return the transformed prompt in case the caller wants to use it i.e. for Blended Create
                return mutablePrompt
            },
            uploadAudio: { audioRecording in
                AsyncThrowingStream<AudioUploadStatus, Error> { continuation in
                    let task = Task {
                        // Create the audio file
                        let uploadRequest: UploadRequest
                        do {
                            uploadRequest = try await apiClient.createAudioFile()
                        } catch {
                            continuation.finish(throwing: AudioUploadError.fileCreationFailed)
                            return
                        }
                        // Upload the audio file while providing the stream with progress
                        for await event in apiClient.uploadStream(input: audioRecording.url, uploadRequest: uploadRequest) {
                            continuation.yield(.uploadEvent(event))
                            // Go to next step after success or handle failure
                            switch event {
                            case .success:
                                break
                            case .failure(let error):
                                continuation.finish(throwing: error)
                                return
                            case .updateProgress:
                                continue
                            }
                        }
                        // Finish the audio upload
                        do {
                            try await apiClient.finishAudioUpload(uploadRequest.id, audioRecording.title)
                        } catch {
                            continuation.finish(throwing: AudioUploadError.verificationFailed)
                            return
                        }

                        // Poll the status of the upload
                        do {
                            for await _ in clock.timer(interval: .seconds(3)) {
                                let audioStatus = try await apiClient.uploadAudioStatus(uploadRequest.id)
                                continuation.yield(.status(audioStatus))
                                if audioStatus.errorMessage != nil || audioStatus.status == .complete {
                                    continuation.finish()
                                    return
                                }
                            }
                        } catch {
                            continuation.finish(throwing: AudioUploadError.statusCheckFailed)
                            return
                        }
                    }

                    continuation.onTermination = { @Sendable _ in
                        task.cancel()
                    }
                }
            },
            initializeClip: { uploadRequestId in
                // TODO: Migrate to APIClientV2
                guard let clipId = try await apiClient.initializeClip(uploadId: uploadRequestId) else { throw InitializeClipError.noClipIdReturned }
                return clipId
            },
            generateMore: { prompt, createType, lyricsModel in
                var mutablePrompt = prompt
                mutablePrompt.lyricsModel = lyricsModel ?? prompt.lyricsModel

                var isAudioUpload: Bool = false

                // Audio create mode transforms
                if case .audio(let clipId) = createType {
                    isAudioUpload = true
                    switch mutablePrompt.audioCreateStyle {
                    case .cover:
                        mutablePrompt.coverClipId = mutablePrompt.clipId
                        mutablePrompt.continueClipId = nil
                        mutablePrompt.continueAt = 0
                        mutablePrompt.task = .cover

                    case .extend:
                        mutablePrompt.continueClipId = mutablePrompt.clipId
                        mutablePrompt.coverClipId = nil
                        mutablePrompt.continueAt = mutablePrompt.audioRecording?.duration
                        mutablePrompt.task = .uploadExtend
                    }
                }

                let model = sunoModelClient.getCurrentModelExternalKey(isAudioUpload: isAudioUpload)

                let clips = try await _generateWithHCaptchaRetries(mutablePrompt, model: model, lyricsModel: lyricsModel)

                @Shared(.fileStorage(.savedPrompts)) var savedPrompts: [Clip.ID: Prompt] = [:]
                $savedPrompts.withLock {
                    for clip in clips {
                        $0[clip.id] = mutablePrompt.toPrompt()
                    }
                }

                return clips
            },
            generateResultStream: {
                UncheckedSendable(subject.values)
                    .eraseToStream()
            }
        )
    }()
}

extension BlendedCreateClient: TestDependencyKey {
    public static let previewValue = Self()
    public static let testValue = Self()
}

public extension DependencyValues {
    var blendedCreateClient: BlendedCreateClient {
        get { self[BlendedCreateClient.self] }
        set { self[BlendedCreateClient.self] = newValue }
    }
}