package io.synor.compute;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import okhttp3.sse.EventSources;

import java.io.Closeable;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

/**
 * Synor Compute SDK - Java Client
 *
 * <p>Access distributed heterogeneous compute resources (CPU, GPU, TPU, NPU, LPU, FPGA, DSP)
 * for AI/ML workloads at 90% cost reduction compared to traditional cloud.</p>
 *
 * <h2>Quick Start</h2>
 * <pre>{@code
 * // Create client
 * SynorCompute client = new SynorCompute("your-api-key");
 *
 * // Matrix multiplication on GPU
 * Tensor a = Tensor.rand(512, 512);
 * Tensor b = Tensor.rand(512, 512);
 * JobResult<Tensor> result = client.matmul(a, b, MatMulOptions.builder()
 *     .processor(ProcessorType.GPU)
 *     .precision(Precision.FP16)
 *     .build());
 *
 * if (result.isSuccess()) {
 *     System.out.println("Result shape: " + Arrays.toString(result.getResult().getShape()));
 *     System.out.println("Time: " + result.getExecutionTimeMs() + "ms");
 * }
 *
 * // LLM inference
 * JobResult<String> response = client.inference("llama-3-70b", "Explain quantum computing");
 * System.out.println(response.getResult());
 *
 * // Streaming inference
 * client.inferenceStream("llama-3-70b", "Write a poem about AI", token -> {
 *     System.out.print(token);
 * });
 *
 * // Clean up
 * client.close();
 * }</pre>
 */
public class SynorCompute implements Closeable {
    private static final MediaType JSON = MediaType.parse("application/json; charset=utf-8");
    private static final String VERSION = "0.1.0";

    private final SynorConfig config;
    private final OkHttpClient httpClient;
    private final ObjectMapper objectMapper;
    private volatile boolean closed = false;

    /**
     * Create a new client with API key.
     */
    public SynorCompute(String apiKey) {
        this(SynorConfig.builder(apiKey).build());
    }

    /**
     * Create a new client with configuration.
     */
    public SynorCompute(SynorConfig config) {
        this.config = config;
        this.httpClient = new OkHttpClient.Builder()
                .connectTimeout(config.getTimeoutMs(), TimeUnit.MILLISECONDS)
                .readTimeout(config.getTimeoutMs(), TimeUnit.MILLISECONDS)
                .writeTimeout(config.getTimeoutMs(), TimeUnit.MILLISECONDS)
                .build();
        this.objectMapper = new ObjectMapper()
                .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
    }

    // ==================== Matrix Operations ====================

    /**
     * Perform matrix multiplication.
     */
    public JobResult<Tensor> matmul(Tensor a, Tensor b) throws IOException {
        return matmul(a, b, MatMulOptions.builder().build());
    }

    /**
     * Perform matrix multiplication with options.
     */
    public JobResult<Tensor> matmul(Tensor a, Tensor b, MatMulOptions options) throws IOException {
        checkNotClosed();

        Map<String, Object> body = new HashMap<>();
        body.put("operation", "matmul");
        body.put("a", tensorToMap(a));
        body.put("b", tensorToMap(b));
        body.put("precision", options.getPrecision().getValue());
        body.put("processor", options.getProcessor().getValue());
        body.put("priority", options.getPriority().getValue());

        Map<String, Object> response = post("/compute", body);
        return parseJobResult(response, Tensor.class);
    }

    /**
     * Perform 2D convolution.
     */
    public JobResult<Tensor> conv2d(Tensor input, Tensor kernel, Conv2dOptions options) throws IOException {
        checkNotClosed();

        Map<String, Object> body = new HashMap<>();
        body.put("operation", "conv2d");
        body.put("input", tensorToMap(input));
        body.put("kernel", tensorToMap(kernel));
        body.put("stride", options.getStride());
        body.put("padding", options.getPadding());
        body.put("precision", options.getPrecision().getValue());

        Map<String, Object> response = post("/compute", body);
        return parseJobResult(response, Tensor.class);
    }

    /**
     * Perform attention computation.
     */
    public JobResult<Tensor> attention(Tensor query, Tensor key, Tensor value,
                                        AttentionOptions options) throws IOException {
        checkNotClosed();

        Map<String, Object> body = new HashMap<>();
        body.put("operation", "attention");
        body.put("query", tensorToMap(query));
        body.put("key", tensorToMap(key));
        body.put("value", tensorToMap(value));
        body.put("num_heads", options.getNumHeads());
        body.put("flash", options.isFlash());
        body.put("precision", options.getPrecision().getValue());

        Map<String, Object> response = post("/compute", body);
        return parseJobResult(response, Tensor.class);
    }

    // ==================== LLM Inference ====================

    /**
     * Run inference on a model.
     */
    public JobResult<String> inference(String model, String prompt) throws IOException {
        return inference(model, prompt, InferenceOptions.builder().build());
    }

    /**
     * Run inference with options.
     */
    public JobResult<String> inference(String model, String prompt, InferenceOptions options)
            throws IOException {
        checkNotClosed();

        Map<String, Object> body = new HashMap<>();
        body.put("operation", "inference");
        body.put("model", model);
        body.put("prompt", prompt);
        body.put("max_tokens", options.getMaxTokens());
        body.put("temperature", options.getTemperature());
        body.put("top_p", options.getTopP());
        body.put("top_k", options.getTopK());
        if (options.getProcessor() != null) {
            body.put("processor", options.getProcessor().getValue());
        }

        Map<String, Object> response = post("/inference", body);
        return parseJobResult(response, String.class);
    }

    /**
     * Run streaming inference.
     */
    public void inferenceStream(String model, String prompt, Consumer<String> onToken)
            throws IOException {
        inferenceStream(model, prompt, InferenceOptions.builder().build(), onToken);
    }

    /**
     * Run streaming inference with options.
     */
    public void inferenceStream(String model, String prompt, InferenceOptions options,
                                 Consumer<String> onToken) throws IOException {
        checkNotClosed();

        Map<String, Object> body = new HashMap<>();
        body.put("operation", "inference");
        body.put("model", model);
        body.put("prompt", prompt);
        body.put("max_tokens", options.getMaxTokens());
        body.put("temperature", options.getTemperature());
        body.put("stream", true);

        Request request = new Request.Builder()
                .url(config.getBaseUrl() + "/inference/stream")
                .addHeader("Authorization", "Bearer " + config.getApiKey())
                .addHeader("Content-Type", "application/json")
                .addHeader("X-SDK-Version", "java/" + VERSION)
                .post(RequestBody.create(objectMapper.writeValueAsString(body), JSON))
                .build();

        CompletableFuture<Void> future = new CompletableFuture<>();

        EventSource.Factory factory = EventSources.createFactory(httpClient);
        factory.newEventSource(request, new EventSourceListener() {
            @Override
            public void onEvent(EventSource source, String id, String type, String data) {
                if ("[DONE]".equals(data)) {
                    future.complete(null);
                    return;
                }
                try {
                    Map<String, Object> json = objectMapper.readValue(data, new TypeReference<>() {});
                    String token = (String) json.get("token");
                    if (token != null) {
                        onToken.accept(token);
                    }
                } catch (JsonProcessingException e) {
                    // Skip malformed JSON
                }
            }

            @Override
            public void onFailure(EventSource source, Throwable t, Response response) {
                future.completeExceptionally(t != null ? t : new IOException("Stream failed"));
            }

            @Override
            public void onClosed(EventSource source) {
                future.complete(null);
            }
        });

        future.join();
    }

    /**
     * Async inference returning CompletableFuture.
     */
    public CompletableFuture<JobResult<String>> inferenceAsync(String model, String prompt) {
        return inferenceAsync(model, prompt, InferenceOptions.builder().build());
    }

    /**
     * Async inference with options.
     */
    public CompletableFuture<JobResult<String>> inferenceAsync(String model, String prompt,
                                                                 InferenceOptions options) {
        return CompletableFuture.supplyAsync(() -> {
            try {
                return inference(model, prompt, options);
            } catch (IOException e) {
                return JobResult.failure(null, e.getMessage());
            }
        });
    }

    // ==================== Model Registry ====================

    /**
     * List available models.
     */
    public List<ModelInfo> listModels() throws IOException {
        return listModels(null);
    }

    /**
     * List models by category.
     */
    public List<ModelInfo> listModels(ModelCategory category) throws IOException {
        checkNotClosed();

        String url = "/models";
        if (category != null) {
            url += "?category=" + category.getValue();
        }

        Map<String, Object> response = get(url);
        List<Map<String, Object>> models = (List<Map<String, Object>>) response.get("models");
        List<ModelInfo> result = new ArrayList<>();
        for (Map<String, Object> m : models) {
            result.add(objectMapper.convertValue(m, ModelInfo.class));
        }
        return result;
    }

    /**
     * Get model by ID.
     */
    public ModelInfo getModel(String modelId) throws IOException {
        checkNotClosed();
        Map<String, Object> response = get("/models/" + modelId);
        return objectMapper.convertValue(response, ModelInfo.class);
    }

    /**
     * Search models.
     */
    public List<ModelInfo> searchModels(String query) throws IOException {
        checkNotClosed();
        Map<String, Object> response = get("/models/search?q=" + query);
        List<Map<String, Object>> models = (List<Map<String, Object>>) response.get("models");
        List<ModelInfo> result = new ArrayList<>();
        for (Map<String, Object> m : models) {
            result.add(objectMapper.convertValue(m, ModelInfo.class));
        }
        return result;
    }

    // ==================== Pricing & Usage ====================

    /**
     * Get current pricing information.
     */
    public List<PricingInfo> getPricing() throws IOException {
        checkNotClosed();
        Map<String, Object> response = get("/pricing");
        List<Map<String, Object>> pricing = (List<Map<String, Object>>) response.get("pricing");
        List<PricingInfo> result = new ArrayList<>();
        for (Map<String, Object> p : pricing) {
            result.add(objectMapper.convertValue(p, PricingInfo.class));
        }
        return result;
    }

    /**
     * Get usage statistics.
     */
    public UsageStats getUsage() throws IOException {
        checkNotClosed();
        Map<String, Object> response = get("/usage");
        return objectMapper.convertValue(response, UsageStats.class);
    }

    // ==================== Health Check ====================

    /**
     * Check service health.
     */
    public boolean healthCheck() {
        try {
            Map<String, Object> response = get("/health");
            return "healthy".equals(response.get("status"));
        } catch (IOException e) {
            return false;
        }
    }

    // ==================== Internal HTTP Methods ====================

    private Map<String, Object> get(String path) throws IOException {
        Request request = new Request.Builder()
                .url(config.getBaseUrl() + path)
                .addHeader("Authorization", "Bearer " + config.getApiKey())
                .addHeader("X-SDK-Version", "java/" + VERSION)
                .get()
                .build();

        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("Request failed: " + response.code());
            }
            return objectMapper.readValue(response.body().string(), new TypeReference<>() {});
        }
    }

    private Map<String, Object> post(String path, Map<String, Object> body) throws IOException {
        Request request = new Request.Builder()
                .url(config.getBaseUrl() + path)
                .addHeader("Authorization", "Bearer " + config.getApiKey())
                .addHeader("Content-Type", "application/json")
                .addHeader("X-SDK-Version", "java/" + VERSION)
                .post(RequestBody.create(objectMapper.writeValueAsString(body), JSON))
                .build();

        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("Request failed: " + response.code());
            }
            return objectMapper.readValue(response.body().string(), new TypeReference<>() {});
        }
    }

    private Map<String, Object> tensorToMap(Tensor tensor) {
        Map<String, Object> map = new HashMap<>();
        map.put("shape", tensor.getShape());
        map.put("data", tensor.getData());
        map.put("dtype", tensor.getDtype().getValue());
        return map;
    }

    @SuppressWarnings("unchecked")
    private <T> JobResult<T> parseJobResult(Map<String, Object> response, Class<T> resultClass)
            throws JsonProcessingException {
        String jobId = (String) response.get("job_id");
        JobStatus status = JobStatus.fromValue((String) response.get("status"));
        String error = (String) response.get("error");
        Long executionTimeMs = response.get("execution_time_ms") != null
                ? ((Number) response.get("execution_time_ms")).longValue()
                : null;
        ProcessorType processor = response.get("processor") != null
                ? ProcessorType.fromValue((String) response.get("processor"))
                : null;
        Double cost = response.get("cost") != null
                ? ((Number) response.get("cost")).doubleValue()
                : null;

        T result = null;
        if (response.get("result") != null) {
            if (resultClass == String.class) {
                result = (T) response.get("result");
            } else if (resultClass == Tensor.class) {
                result = objectMapper.convertValue(response.get("result"), (Class<T>) Tensor.class);
            } else {
                result = objectMapper.convertValue(response.get("result"), resultClass);
            }
        }

        return new JobResult<>(jobId, status, result, error, executionTimeMs, processor, cost);
    }

    private void checkNotClosed() {
        if (closed) {
            throw new IllegalStateException("Client has been closed");
        }
    }

    @Override
    public void close() {
        closed = true;
        httpClient.dispatcher().executorService().shutdown();
        httpClient.connectionPool().evictAll();
    }
}
