import Foundation
import UIKit
@preconcurrency import Alamofire

final class ComfyClient {
    static let shared = ComfyClient()

    private let endpoint = "http://192.168.x.x:8000/dream-weaver"
    private let session: Session

    private init(session: Session = .default) {
        self.session = session
    }

    func generateImage(source: UIImage, prompt: String) async throws -> UIImage {
        let resized = source.resizedSquare(to: 1024)
        guard let imageData = resized.jpegData(compressionQuality: 0.9) else {
            throw ComfyClientError.encodingFailed
        }

        let payload = DreamWeaverRequest(
            imageBase64: imageData.base64EncodedString(),
            prompt: prompt
        )

        let response = try await session.request(
            endpoint,
            method: .post,
            parameters: payload,
            encoder: JSONParameterEncoder.default,
            headers: [.contentType("application/json")]
        )
        .validate(statusCode: 200..<300)
        .serializingDecodable(DreamWeaverResponse.self)
        .value

        guard
            let data = Data(base64Encoded: response.outputBase64),
            let generated = UIImage(data: data)
        else {
            throw ComfyClientError.decodingFailed
        }

        return generated
    }
}

private struct DreamWeaverRequest: Encodable, Sendable {
    let imageBase64: String
    let prompt: String
}

private struct DreamWeaverResponse: Decodable, Sendable {
    let outputBase64: String

    enum CodingKeys: String, CodingKey {
        case outputBase64 = "output_base64"
        case imageBase64 = "image_base64"
        case image
    }

    init(from decoder: Decoder) throws {
        let container = try decoder.container(keyedBy: CodingKeys.self)
        if let preferred = try container.decodeIfPresent(String.self, forKey: .outputBase64) {
            outputBase64 = preferred
            return
        }
        if let legacy = try container.decodeIfPresent(String.self, forKey: .imageBase64) {
            outputBase64 = legacy
            return
        }
        outputBase64 = try container.decode(String.self, forKey: .image)
    }
}

enum ComfyClientError: Error {
    case encodingFailed
    case decodingFailed
}

private extension UIImage {
    func resizedSquare(to side: CGFloat) -> UIImage {
        let format = UIGraphicsImageRendererFormat.default()
        format.scale = 1
        let renderer = UIGraphicsImageRenderer(size: CGSize(width: side, height: side), format: format)

        return renderer.image { _ in
            let aspect = size.width / size.height
            let targetRect: CGRect
            if aspect > 1 {
                let width = side * aspect
                targetRect = CGRect(x: (side - width) / 2, y: 0, width: width, height: side)
            } else {
                let height = side / aspect
                targetRect = CGRect(x: 0, y: (side - height) / 2, width: side, height: height)
            }
            draw(in: targetRect)
        }
    }
}