import Foundation

/// Synor Compute SDK - Swift Client
///
/// Access distributed heterogeneous compute resources (CPU, GPU, TPU, NPU, LPU, FPGA, DSP)
/// for AI/ML workloads at 90% cost reduction compared to traditional cloud.
///
/// ```swift
/// // Create client
/// let client = SynorCompute(apiKey: "your-api-key")
///
/// // Matrix multiplication on GPU
/// let a = Tensor.rand([512, 512])
/// let b = Tensor.rand([512, 512])
/// let result = try await client.matmul(a, b, options: MatMulOptions(
///     processor: .gpu,
///     precision: .fp16
/// ))
///
/// if result.isSuccess {
///     print("Result shape: \(result.result?.shape ?? [])")
///     print("Time: \(result.executionTimeMs ?? 0)ms")
/// }
///
/// // LLM inference
/// let response = try await client.inference("llama-3-70b", prompt: "Explain quantum computing")
/// print(response.result ?? "")
///
/// // Streaming inference
/// for try await token in client.inferenceStream("llama-3-70b", prompt: "Write a poem about AI") {
///     print(token, terminator: "")
/// }
/// ```
public final class SynorCompute {
    public static let version = "0.1.0"

    private let config: SynorConfig
    private let session: URLSession
    private let encoder = JSONEncoder()
    private let decoder = JSONDecoder()
    private var isClosed = false

    public init(apiKey: String) {
        self.config = SynorConfig(apiKey: apiKey)
        self.session = URLSession(configuration: .default)
    }

    public init(config: SynorConfig) {
        self.config = config
        let configuration = URLSessionConfiguration.default
        configuration.timeoutIntervalForRequest = config.timeoutSeconds
        self.session = URLSession(configuration: configuration)
    }

    // MARK: - Matrix Operations

    /// Perform matrix multiplication
    public func matmul(
        _ a: Tensor,
        _ b: Tensor,
        options: MatMulOptions = MatMulOptions()
    ) async throws -> JobResult<Tensor> {
        try checkNotClosed()

        let body: [String: Any] = [
            "operation": "matmul",
            "a": tensorToDict(a),
            "b": tensorToDict(b),
            "precision": options.precision.rawValue,
            "processor": options.processor.rawValue,
            "priority": options.priority.rawValue
        ]

        return try await post("/compute", body: body)
    }

    /// Perform 2D convolution
    public func conv2d(
        _ input: Tensor,
        kernel: Tensor,
        options: Conv2dOptions = Conv2dOptions()
    ) async throws -> JobResult<Tensor> {
        try checkNotClosed()

        let body: [String: Any] = [
            "operation": "conv2d",
            "input": tensorToDict(input),
            "kernel": tensorToDict(kernel),
            "stride": [options.stride.0, options.stride.1],
            "padding": [options.padding.0, options.padding.1],
            "precision": options.precision.rawValue
        ]

        return try await post("/compute", body: body)
    }

    /// Perform attention computation
    public func attention(
        query: Tensor,
        key: Tensor,
        value: Tensor,
        options: AttentionOptions = AttentionOptions()
    ) async throws -> JobResult<Tensor> {
        try checkNotClosed()

        let body: [String: Any] = [
            "operation": "attention",
            "query": tensorToDict(query),
            "key": tensorToDict(key),
            "value": tensorToDict(value),
            "num_heads": options.numHeads,
            "flash": options.flash,
            "precision": options.precision.rawValue
        ]

        return try await post("/compute", body: body)
    }

    // MARK: - LLM Inference

    /// Run inference on a model
    public func inference(
        _ model: String,
        prompt: String,
        options: InferenceOptions = InferenceOptions()
    ) async throws -> JobResult<String> {
        try checkNotClosed()

        var body: [String: Any] = [
            "operation": "inference",
            "model": model,
            "prompt": prompt,
            "max_tokens": options.maxTokens,
            "temperature": options.temperature,
            "top_p": options.topP,
            "top_k": options.topK
        ]

        if let processor = options.processor {
            body["processor"] = processor.rawValue
        }

        return try await postString("/inference", body: body)
    }

    /// Run streaming inference
    public func inferenceStream(
        _ model: String,
        prompt: String,
        options: InferenceOptions = InferenceOptions()
    ) -> AsyncThrowingStream<String, Error> {
        AsyncThrowingStream { continuation in
            Task {
                do {
                    try checkNotClosed()

                    let body: [String: Any] = [
                        "operation": "inference",
                        "model": model,
                        "prompt": prompt,
                        "max_tokens": options.maxTokens,
                        "temperature": options.temperature,
                        "stream": true
                    ]

                    var request = try createRequest("/inference/stream", method: "POST")
                    request.httpBody = try JSONSerialization.data(withJSONObject: body)

                    let (bytes, _) = try await session.bytes(for: request)

                    for try await line in bytes.lines {
                        if line.hasPrefix("data: ") {
                            let data = String(line.dropFirst(6))
                            if data == "[DONE]" {
                                break
                            }
                            if let jsonData = data.data(using: .utf8),
                               let json = try? JSONSerialization.jsonObject(with: jsonData) as? [String: Any],
                               let token = json["token"] as? String {
                                continuation.yield(token)
                            }
                        }
                    }

                    continuation.finish()
                } catch {
                    continuation.finish(throwing: error)
                }
            }
        }
    }

    // MARK: - Model Registry

    /// List available models
    public func listModels(category: ModelCategory? = nil) async throws -> [ModelInfo] {
        try checkNotClosed()

        let path = category.map { "/models?category=\($0.rawValue)" } ?? "/models"
        let response: [String: Any] = try await get(path)
        guard let models = response["models"] as? [[String: Any]] else {
            return []
        }

        return try models.compactMap { dict in
            let data = try JSONSerialization.data(withJSONObject: dict)
            return try decoder.decode(ModelInfo.self, from: data)
        }
    }

    /// Get model by ID
    public func getModel(_ modelId: String) async throws -> ModelInfo {
        try checkNotClosed()
        return try await get("/models/\(modelId)")
    }

    /// Search models
    public func searchModels(_ query: String) async throws -> [ModelInfo] {
        try checkNotClosed()

        let encoded = query.addingPercentEncoding(withAllowedCharacters: .urlQueryAllowed) ?? query
        let response: [String: Any] = try await get("/models/search?q=\(encoded)")
        guard let models = response["models"] as? [[String: Any]] else {
            return []
        }

        return try models.compactMap { dict in
            let data = try JSONSerialization.data(withJSONObject: dict)
            return try decoder.decode(ModelInfo.self, from: data)
        }
    }

    // MARK: - Pricing & Usage

    /// Get current pricing information
    public func getPricing() async throws -> [PricingInfo] {
        try checkNotClosed()

        let response: [String: Any] = try await get("/pricing")
        guard let pricing = response["pricing"] as? [[String: Any]] else {
            return []
        }

        return try pricing.compactMap { dict in
            let data = try JSONSerialization.data(withJSONObject: dict)
            return try decoder.decode(PricingInfo.self, from: data)
        }
    }

    /// Get usage statistics
    public func getUsage() async throws -> UsageStats {
        try checkNotClosed()
        return try await get("/usage")
    }

    // MARK: - Health Check

    /// Check service health
    public func healthCheck() async -> Bool {
        do {
            let response: [String: Any] = try await get("/health")
            return response["status"] as? String == "healthy"
        } catch {
            return false
        }
    }

    // MARK: - Lifecycle

    /// Close the client
    public func close() {
        isClosed = true
        session.invalidateAndCancel()
    }

    // MARK: - Private Methods

    private func checkNotClosed() throws {
        guard !isClosed else {
            throw SynorError.clientClosed
        }
    }

    private func createRequest(_ path: String, method: String) throws -> URLRequest {
        guard let url = URL(string: config.baseUrl + path) else {
            throw SynorError.invalidConfiguration("Invalid URL")
        }

        var request = URLRequest(url: url)
        request.httpMethod = method
        request.setValue("Bearer \(config.apiKey)", forHTTPHeaderField: "Authorization")
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.setValue("swift/\(Self.version)", forHTTPHeaderField: "X-SDK-Version")
        return request
    }

    private func get<T: Decodable>(_ path: String) async throws -> T {
        let request = try createRequest(path, method: "GET")
        let (data, response) = try await session.data(for: request)

        guard let httpResponse = response as? HTTPURLResponse else {
            throw SynorError.networkError(URLError(.badServerResponse))
        }

        guard httpResponse.statusCode == 200 else {
            let message = String(data: data, encoding: .utf8) ?? "Unknown error"
            throw SynorError.apiError(httpResponse.statusCode, message)
        }

        return try decoder.decode(T.self, from: data)
    }

    private func post<T: Decodable>(_ path: String, body: [String: Any]) async throws -> T {
        var request = try createRequest(path, method: "POST")
        request.httpBody = try JSONSerialization.data(withJSONObject: body)

        let (data, response) = try await session.data(for: request)

        guard let httpResponse = response as? HTTPURLResponse else {
            throw SynorError.networkError(URLError(.badServerResponse))
        }

        guard httpResponse.statusCode == 200 else {
            let message = String(data: data, encoding: .utf8) ?? "Unknown error"
            throw SynorError.apiError(httpResponse.statusCode, message)
        }

        return try decoder.decode(T.self, from: data)
    }

    private func postString(_ path: String, body: [String: Any]) async throws -> JobResult<String> {
        var request = try createRequest(path, method: "POST")
        request.httpBody = try JSONSerialization.data(withJSONObject: body)

        let (data, response) = try await session.data(for: request)

        guard let httpResponse = response as? HTTPURLResponse else {
            throw SynorError.networkError(URLError(.badServerResponse))
        }

        guard httpResponse.statusCode == 200 else {
            let message = String(data: data, encoding: .utf8) ?? "Unknown error"
            throw SynorError.apiError(httpResponse.statusCode, message)
        }

        guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            throw SynorError.decodingError(DecodingError.dataCorrupted(.init(codingPath: [], debugDescription: "Invalid JSON")))
        }

        return JobResult(
            jobId: json["job_id"] as? String,
            status: (json["status"] as? String).flatMap { JobStatus(rawValue: $0) } ?? .pending,
            result: json["result"] as? String,
            error: json["error"] as? String,
            executionTimeMs: json["execution_time_ms"] as? Int64,
            processor: (json["processor"] as? String).flatMap { ProcessorType(rawValue: $0) },
            cost: json["cost"] as? Double
        )
    }

    private func tensorToDict(_ tensor: Tensor) -> [String: Any] {
        [
            "shape": tensor.shape,
            "data": tensor.data,
            "dtype": tensor.dtype.rawValue
        ]
    }
}