import Foundation

// Reference transport for a native app, not an official Antla native SDK.
// Persist visitor/job IDs in your app and keep the SDK key out of source control.
struct AntlaClient {
    let baseURL: URL // Includes /api/headless
    let sdkKey: String
    var session: URLSession = .shared

    struct Envelope<T: Decodable>: Decodable { let data: T }
    struct Failure: Decodable {
        let error: String
        let code: String?
        let message: String?
    }
    struct APIError: Error {
        let status: Int
        let code: String?
        let message: String
        let retryAfter: Double?
    }
    struct SubmissionUncertain: Error {}
    struct ProtocolError: Error {}
    struct PollingTimeout: Error {}

    struct Upload: Decodable { let uploadUrl: URL; let imageUrl: URL }
    struct Visitor: Decodable {
        let uniqueUserId: String?
        let requestsLeft: Int
        let totalRequests: Int
    }
    struct StartInput: Encodable {
        let funnelId: Int
        let productId: String
        let variantId: String
        let imageUrl: URL
        var uniqueUserId: String?
        var customerEmail: String?
        var moreLooksConsent: Bool?
    }
    struct Started: Decodable {
        let jobId: String
        let etaSeconds: Int?
        let uniqueUserId: String?
    }
    struct Generation: Decodable {
        let status: String
        let storageStatus: String?
        let imageUrl: URL?
        let displayUrl: URL?
        let permanentUrl: URL?
        let error: String?
        let uniqueUserId: String?
    }

    private func request<T: Decodable>(
        _ path: String, method: String = "GET",
        body: Data? = nil, visitor: String? = nil
    ) async throws -> T {
        var components = URLComponents(
            url: baseURL.appendingPathComponent(path),
            resolvingAgainstBaseURL: false
        )!
        if let visitor, !visitor.isEmpty {
            components.queryItems = [.init(name: "uniqueUserId", value: visitor)]
        }
        guard let url = components.url, url.scheme == "https" else { throw ProtocolError() }
        var req = URLRequest(url: url)
        req.httpMethod = method
        req.timeoutInterval = 20
        req.setValue(sdkKey, forHTTPHeaderField: "X-Antla-SDK-Key")
        req.httpBody = body
        if body != nil { req.setValue("application/json", forHTTPHeaderField: "Content-Type") }
        let (bytes, response) = try await session.data(for: req)
        guard let http = response as? HTTPURLResponse else { throw ProtocolError() }
        guard (200..<300).contains(http.statusCode) else {
            let failure = try? JSONDecoder().decode(Failure.self, from: bytes)
            throw APIError(
                status: http.statusCode, code: failure?.code,
                message: failure?.message ?? failure?.error ?? "Antla request failed",
                retryAfter: http.value(forHTTPHeaderField: "Retry-After").flatMap(Double.init)
            )
        }
        return try JSONDecoder().decode(Envelope<T>.self, from: bytes).data
    }

    func createUpload(contentType: String) async throws -> Upload {
        try await request("uploads", method: "POST",
            body: JSONSerialization.data(withJSONObject: ["contentType": contentType]))
    }

    func upload(_ bytes: Data, to url: URL, contentType: String) async throws {
        guard url.scheme == "https" else { throw ProtocolError() }
        var req = URLRequest(url: url)
        req.httpMethod = "PUT"
        req.timeoutInterval = 60
        req.httpBody = bytes
        req.setValue(contentType, forHTTPHeaderField: "Content-Type")
        // Do not add the Antla SDK key to this storage request.
        let (_, response) = try await session.data(for: req)
        guard let http = response as? HTTPURLResponse,
              (200..<300).contains(http.statusCode) else { throw ProtocolError() }
    }

    func requestsLeft(visitor: String?) async throws -> Visitor {
        try await request("requests-left", visitor: visitor)
    }

    func register(email: String, name: String? = nil, visitor: String? = nil) async throws -> Visitor {
        var body = ["email": email]
        if let name { body["name"] = name }
        if let visitor { body["uniqueUserId"] = visitor }
        return try await request("customer-register", method: "POST",
            body: JSONSerialization.data(withJSONObject: body))
    }

    func start(_ input: StartInput) async throws -> Started {
        let body = try JSONEncoder().encode(input)
        do {
            let started: Started = try await request("generations", method: "POST", body: body)
            guard !started.jobId.isEmpty else { throw ProtocolError() }
            return started
        } catch let error as APIError where (400..<500).contains(error.status) {
            throw error
        } catch {
            // Never automatically replay this POST.
            throw SubmissionUncertain()
        }
    }

    func status(jobID: String, visitor: String?) async throws -> Generation {
        guard UUID(uuidString: jobID) != nil else { throw ProtocolError() }
        return try await request("generations/" + jobID, visitor: visitor)
    }

    func poll(
        jobID: String, visitor: String?,
        timeout: TimeInterval = 180,
        onPreview: ((URL) -> Void)? = nil,
        onIdentity: ((String) -> Void)? = nil
    ) async throws -> URL {
        let deadline = Date().addingTimeInterval(timeout)
        var visitorID = visitor
        var failures = 0
        while Date() < deadline {
            try Task.checkCancellation()
            let result: Generation
            do {
                result = try await status(jobID: jobID, visitor: visitorID)
                failures = 0
            } catch {
                try Task.checkCancellation()
                failures += 1
                guard failures <= 3 else { throw error }
                let wait: Double
                if let api = error as? APIError {
                    if api.status == 429, api.code != "TRY_ON_LIMIT_REACHED",
                       let retry = api.retryAfter, retry.isFinite, retry >= 0 {
                        wait = max(1, retry)
                    } else if api.status >= 500 {
                        wait = min(8, pow(2, Double(failures)))
                    } else { throw error }
                } else if error is URLError {
                    wait = min(8, pow(2, Double(failures)))
                } else { throw error }
                guard Date().addingTimeInterval(wait) < deadline else { throw PollingTimeout() }
                try await Task.sleep(nanoseconds: UInt64(wait * 1_000_000_000))
                continue
            }
            if let id = result.uniqueUserId {
                visitorID = id
                onIdentity?(id)
            }
            if result.status == "ERROR" {
                throw APIError(status: 200, code: nil,
                    message: result.error ?? "Generation failed", retryAfter: nil)
            }
            if result.status == "SUCCESS" {
                if result.storageStatus == "PENDING" {
                    if let preview = result.displayUrl ?? result.imageUrl { onPreview?(preview) }
                } else if result.storageStatus == "READY" || result.storageStatus == nil {
                    guard let image = result.permanentUrl ?? result.imageUrl else { throw ProtocolError() }
                    return image
                } else { throw ProtocolError() }
            } else if result.status != "PROCESSING" { throw ProtocolError() }
            guard Date().addingTimeInterval(2) < deadline else { throw PollingTimeout() }
            try await Task.sleep(nanoseconds: 2_000_000_000)
        }
        throw PollingTimeout()
    }
}