package io.synor.zk;

import io.synor.zk.Types.*;

import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.reflect.TypeToken;

import java.io.IOException;
import java.lang.reflect.Type;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Synor ZK SDK Client for Java
 *
 * Zero-Knowledge proof systems for ZK-Rollups and privacy.
 */
public class SynorZk {
    private final ZkConfig config;
    private final HttpClient httpClient;
    private final Gson gson;
    private final AtomicBoolean closed = new AtomicBoolean(false);

    public final ProofsClient proofs;
    public final CircuitsClient circuits;
    public final RollupClient rollup;
    public final StateClient state;
    public final CeremonyClient ceremony;

    public SynorZk(ZkConfig config) {
        this.config = config;
        this.httpClient = HttpClient.newBuilder()
                .connectTimeout(Duration.ofMillis(config.getTimeout()))
                .build();
        this.gson = new GsonBuilder().create();

        this.proofs = new ProofsClient(this);
        this.circuits = new CircuitsClient(this);
        this.rollup = new RollupClient(this);
        this.state = new StateClient(this);
        this.ceremony = new CeremonyClient(this);
    }

    public ProofSystem getDefaultProofSystem() {
        return config.getDefaultProofSystem();
    }

    public CompletableFuture<Boolean> healthCheck() {
        return get("/health", Map.class)
                .thenApply(result -> "healthy".equals(result.get("status")))
                .exceptionally(e -> false);
    }

    public CompletableFuture<Map<String, Object>> getInfo() {
        return get("/info", new TypeToken<Map<String, Object>>(){}.getType());
    }

    public void close() {
        closed.set(true);
    }

    public boolean isClosed() {
        return closed.get();
    }

    // Internal HTTP methods
    <T> CompletableFuture<T> get(String path, Type type) {
        return request("GET", path, null, type);
    }

    <T> CompletableFuture<T> post(String path, Object body, Type type) {
        return request("POST", path, body, type);
    }

    private <T> CompletableFuture<T> request(String method, String path, Object body, Type type) {
        if (closed.get()) {
            return CompletableFuture.failedFuture(
                    new ZkException("Client has been closed", "CLIENT_CLOSED"));
        }

        String url = config.getEndpoint() + path;
        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
                .uri(URI.create(url))
                .header("Content-Type", "application/json")
                .header("Authorization", "Bearer " + config.getApiKey())
                .header("X-SDK-Version", "java/0.1.0")
                .timeout(Duration.ofMillis(config.getTimeout()));

        if ("GET".equals(method)) {
            requestBuilder.GET();
        } else if ("POST".equals(method)) {
            String jsonBody = body != null ? gson.toJson(body) : "{}";
            requestBuilder.POST(HttpRequest.BodyPublishers.ofString(jsonBody));
        } else if ("DELETE".equals(method)) {
            requestBuilder.DELETE();
        }

        return httpClient.sendAsync(requestBuilder.build(), HttpResponse.BodyHandlers.ofString())
                .thenApply(response -> {
                    if (response.statusCode() >= 400) {
                        Map<String, Object> error = gson.fromJson(response.body(), Map.class);
                        throw new ZkException(
                                (String) error.getOrDefault("message", "HTTP " + response.statusCode()),
                                (String) error.get("code"),
                                response.statusCode()
                        );
                    }
                    return gson.fromJson(response.body(), type);
                });
    }

    /** Proofs sub-client */
    public static class ProofsClient {
        private final SynorZk zk;

        ProofsClient(SynorZk zk) {
            this.zk = zk;
        }

        public CompletableFuture<Proof> generate(ProofRequest request) {
            Map<String, Object> body = new HashMap<>();
            body.put("circuit_type", request.getCircuitType().getValue());
            body.put("public_inputs", request.getPublicInputs());
            body.put("private_inputs", request.getPrivateInputs());
            body.put("system", (request.getSystem() != null ? request.getSystem() : zk.getDefaultProofSystem()).getValue());
            return zk.post("/proofs/generate", body, Proof.class);
        }

        public CompletableFuture<VerificationResult> verify(Proof proof) {
            Map<String, Object> body = new HashMap<>();
            body.put("proof_id", proof.getId());
            body.put("data", proof.getData());
            body.put("public_inputs", proof.getPublicInputs());
            body.put("system", proof.getSystem().getValue());
            body.put("circuit_type", proof.getCircuitType().getValue());
            return zk.post("/proofs/verify", body, VerificationResult.class);
        }

        public CompletableFuture<Proof> get(String proofId) {
            return zk.get("/proofs/" + proofId, Proof.class);
        }

        public CompletableFuture<List<Proof>> list(Integer limit, Integer offset) {
            StringBuilder path = new StringBuilder("/proofs");
            if (limit != null || offset != null) {
                path.append("?");
                if (limit != null) path.append("limit=").append(limit);
                if (limit != null && offset != null) path.append("&");
                if (offset != null) path.append("offset=").append(offset);
            }
            return zk.get(path.toString(), new TypeToken<Map<String, List<Proof>>>(){}.getType())
                    .thenApply(result -> ((Map<String, List<Proof>>) result).getOrDefault("proofs", Collections.emptyList()));
        }

        public CompletableFuture<byte[]> serialize(Proof proof) {
            Map<String, Object> body = new HashMap<>();
            body.put("proof_id", proof.getId());
            return zk.post("/proofs/serialize", body, new TypeToken<Map<String, String>>(){}.getType())
                    .thenApply(result -> Base64.getDecoder().decode(((Map<String, String>) result).get("data")));
        }

        public CompletableFuture<Proof> deserialize(byte[] data, ProofSystem system) {
            Map<String, Object> body = new HashMap<>();
            body.put("data", Base64.getEncoder().encodeToString(data));
            body.put("system", system.getValue());
            return zk.post("/proofs/deserialize", body, Proof.class);
        }
    }

    /** Circuits sub-client */
    public static class CircuitsClient {
        private final SynorZk zk;

        CircuitsClient(SynorZk zk) {
            this.zk = zk;
        }

        public CompletableFuture<VerificationKey> getVerificationKey(CircuitType circuitType, ProofSystem system) {
            ProofSystem sys = system != null ? system : zk.getDefaultProofSystem();
            return zk.get("/circuits/" + circuitType.getValue() + "/vk?system=" + sys.getValue(), VerificationKey.class);
        }

        public CompletableFuture<ProvingKey> getProvingKey(CircuitType circuitType, ProofSystem system) {
            ProofSystem sys = system != null ? system : zk.getDefaultProofSystem();
            return zk.get("/circuits/" + circuitType.getValue() + "/pk?system=" + sys.getValue(), ProvingKey.class);
        }

        public CompletableFuture<List<CircuitConfig>> list() {
            return zk.get("/circuits", new TypeToken<Map<String, List<CircuitConfig>>>(){}.getType())
                    .thenApply(result -> ((Map<String, List<CircuitConfig>>) result).getOrDefault("circuits", Collections.emptyList()));
        }

        public CompletableFuture<CircuitConfig> getConfig(CircuitType circuitType) {
            return zk.get("/circuits/" + circuitType.getValue() + "/config", CircuitConfig.class);
        }

        public CompletableFuture<String> compile(CircuitConfig config) {
            return zk.post("/circuits/compile", config, new TypeToken<Map<String, String>>(){}.getType())
                    .thenApply(result -> ((Map<String, String>) result).get("circuit_id"));
        }
    }

    /** Rollup sub-client */
    public static class RollupClient {
        private final SynorZk zk;

        RollupClient(SynorZk zk) {
            this.zk = zk;
        }

        public CompletableFuture<RollupStats> getStats() {
            return zk.get("/rollup/stats", RollupStats.class);
        }

        public CompletableFuture<Batch> getCurrentBatch() {
            return zk.get("/rollup/batch/current", Batch.class);
        }

        public CompletableFuture<Batch> getBatch(long batchNumber) {
            return zk.get("/rollup/batch/" + batchNumber, Batch.class);
        }

        public CompletableFuture<String> submitTransfer(Transfer transfer) {
            return zk.post("/rollup/transfer", transfer, new TypeToken<Map<String, String>>(){}.getType())
                    .thenApply(result -> ((Map<String, String>) result).get("tx_id"));
        }

        public CompletableFuture<String> submitDeposit(Deposit deposit) {
            return zk.post("/rollup/deposit", deposit, new TypeToken<Map<String, String>>(){}.getType())
                    .thenApply(result -> ((Map<String, String>) result).get("tx_id"));
        }

        public CompletableFuture<String> submitWithdrawal(Withdrawal withdrawal) {
            return zk.post("/rollup/withdraw", withdrawal, new TypeToken<Map<String, String>>(){}.getType())
                    .thenApply(result -> ((Map<String, String>) result).get("tx_id"));
        }

        public CompletableFuture<Batch> finalizeBatch() {
            return zk.post("/rollup/batch/finalize", Collections.emptyMap(), Batch.class);
        }

        public CompletableFuture<List<Transfer>> getPendingTransactions() {
            return zk.get("/rollup/pending", new TypeToken<Map<String, List<Transfer>>>(){}.getType())
                    .thenApply(result -> ((Map<String, List<Transfer>>) result).getOrDefault("transactions", Collections.emptyList()));
        }
    }

    /** State sub-client */
    public static class StateClient {
        private final SynorZk zk;

        StateClient(SynorZk zk) {
            this.zk = zk;
        }

        public CompletableFuture<String> getRoot() {
            return zk.get("/state/root", new TypeToken<Map<String, String>>(){}.getType())
                    .thenApply(result -> ((Map<String, String>) result).get("root"));
        }

        public CompletableFuture<AccountState> getAccount(String address) {
            return zk.get("/state/account/" + address, AccountState.class);
        }

        public CompletableFuture<MerkleProof> getMerkleProof(String address) {
            return zk.get("/state/proof/" + address, MerkleProof.class);
        }

        public CompletableFuture<Boolean> verifyMerkleProof(MerkleProof proof) {
            return zk.post("/state/proof/verify", proof, new TypeToken<Map<String, Boolean>>(){}.getType())
                    .thenApply(result -> ((Map<String, Boolean>) result).get("valid"));
        }

        public CompletableFuture<Map<String, Object>> getStateAtBatch(long batchNumber) {
            return zk.get("/state/batch/" + batchNumber, new TypeToken<Map<String, Object>>(){}.getType());
        }
    }

    /** Ceremony sub-client */
    public static class CeremonyClient {
        private final SynorZk zk;

        CeremonyClient(SynorZk zk) {
            this.zk = zk;
        }

        public CompletableFuture<TrustedSetup> getStatus(CircuitType circuitType, ProofSystem system) {
            ProofSystem sys = system != null ? system : zk.getDefaultProofSystem();
            return zk.get("/ceremony/" + circuitType.getValue() + "?system=" + sys.getValue(), TrustedSetup.class);
        }

        public CompletableFuture<CeremonyContribution> contribute(CircuitType circuitType, byte[] entropy, ProofSystem system) {
            ProofSystem sys = system != null ? system : zk.getDefaultProofSystem();
            Map<String, Object> body = new HashMap<>();
            body.put("circuit_type", circuitType.getValue());
            body.put("entropy", Base64.getEncoder().encodeToString(entropy));
            body.put("system", sys.getValue());
            return zk.post("/ceremony/contribute", body, CeremonyContribution.class);
        }

        public CompletableFuture<Boolean> verifyContribution(CircuitType circuitType, String contributionHash) {
            Map<String, Object> body = new HashMap<>();
            body.put("circuit_type", circuitType.getValue());
            body.put("contribution_hash", contributionHash);
            return zk.post("/ceremony/verify", body, new TypeToken<Map<String, Boolean>>(){}.getType())
                    .thenApply(result -> ((Map<String, Boolean>) result).get("valid"));
        }

        public CompletableFuture<List<CeremonyContribution>> listContributions(CircuitType circuitType) {
            return zk.get("/ceremony/" + circuitType.getValue() + "/contributions", new TypeToken<Map<String, List<CeremonyContribution>>>(){}.getType())
                    .thenApply(result -> ((Map<String, List<CeremonyContribution>>) result).getOrDefault("contributions", Collections.emptyList()));
        }
    }
}
