// Package zk provides the Synor ZK SDK for Go.
package zk

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sync/atomic"
	"time"
)

// Client is the main ZK SDK client
type Client struct {
	config     Config
	httpClient *http.Client
	closed     atomic.Bool

	Proofs   *ProofsClient
	Circuits *CircuitsClient
	Rollup   *RollupClient
	State    *StateClient
	Ceremony *CeremonyClient
}

// New creates a new ZK SDK client
func New(config Config) *Client {
	c := &Client{
		config: config,
		httpClient: &http.Client{
			Timeout: config.Timeout,
		},
	}

	c.Proofs = &ProofsClient{client: c}
	c.Circuits = &CircuitsClient{client: c}
	c.Rollup = &RollupClient{client: c}
	c.State = &StateClient{client: c}
	c.Ceremony = &CeremonyClient{client: c}

	return c
}

// DefaultProofSystem returns the default proof system
func (c *Client) DefaultProofSystem() ProofSystem {
	return c.config.DefaultProofSystem
}

// HealthCheck checks if the service is healthy
func (c *Client) HealthCheck(ctx context.Context) (bool, error) {
	var result map[string]interface{}
	if err := c.get(ctx, "/health", &result); err != nil {
		return false, nil
	}
	return result["status"] == "healthy", nil
}

// GetInfo returns service info
func (c *Client) GetInfo(ctx context.Context) (map[string]interface{}, error) {
	var result map[string]interface{}
	err := c.get(ctx, "/info", &result)
	return result, err
}

// Close closes the client
func (c *Client) Close() {
	c.closed.Store(true)
}

// IsClosed returns whether the client is closed
func (c *Client) IsClosed() bool {
	return c.closed.Load()
}

// Internal HTTP methods
func (c *Client) get(ctx context.Context, path string, result interface{}) error {
	return c.request(ctx, "GET", path, nil, result)
}

func (c *Client) post(ctx context.Context, path string, body, result interface{}) error {
	return c.request(ctx, "POST", path, body, result)
}

func (c *Client) delete(ctx context.Context, path string, result interface{}) error {
	return c.request(ctx, "DELETE", path, nil, result)
}

func (c *Client) request(ctx context.Context, method, path string, body, result interface{}) error {
	if c.closed.Load() {
		return &Error{Message: "Client has been closed", Code: "CLIENT_CLOSED"}
	}

	url := c.config.Endpoint + path

	var bodyReader io.Reader
	if body != nil {
		data, err := json.Marshal(body)
		if err != nil {
			return err
		}
		bodyReader = bytes.NewReader(data)
	}

	var lastErr error
	for attempt := 0; attempt <= c.config.Retries; attempt++ {
		req, err := http.NewRequestWithContext(ctx, method, url, bodyReader)
		if err != nil {
			return err
		}

		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Bearer "+c.config.APIKey)
		req.Header.Set("X-SDK-Version", "go/0.1.0")

		resp, err := c.httpClient.Do(req)
		if err != nil {
			lastErr = err
			if attempt < c.config.Retries {
				time.Sleep(time.Duration(100*(1<<attempt)) * time.Millisecond)
				continue
			}
			break
		}
		defer resp.Body.Close()

		respBody, err := io.ReadAll(resp.Body)
		if err != nil {
			return err
		}

		if resp.StatusCode >= 400 {
			var errResp struct {
				Message string `json:"message"`
				Code    string `json:"code"`
			}
			json.Unmarshal(respBody, &errResp)
			return &Error{
				Message: errResp.Message,
				Code:    errResp.Code,
				Status:  resp.StatusCode,
			}
		}

		if result != nil {
			return json.Unmarshal(respBody, result)
		}
		return nil
	}

	return &Error{
		Message: fmt.Sprintf("request failed: %v", lastErr),
		Code:    "NETWORK_ERROR",
	}
}

// ProofsClient handles proof operations
type ProofsClient struct {
	client *Client
}

// Generate generates a zero-knowledge proof
func (p *ProofsClient) Generate(ctx context.Context, req ProofRequest) (*Proof, error) {
	system := req.System
	if system == nil {
		s := p.client.DefaultProofSystem()
		system = &s
	}

	body := map[string]interface{}{
		"circuit_type":   req.CircuitType,
		"public_inputs":  req.PublicInputs,
		"private_inputs": req.PrivateInputs,
		"system":         *system,
	}

	var result Proof
	if err := p.client.post(ctx, "/proofs/generate", body, &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// Verify verifies a proof
func (p *ProofsClient) Verify(ctx context.Context, proof *Proof) (*VerificationResult, error) {
	body := map[string]interface{}{
		"proof_id":      proof.ID,
		"data":          proof.Data,
		"public_inputs": proof.PublicInputs,
		"system":        proof.System,
		"circuit_type":  proof.CircuitType,
	}

	var result VerificationResult
	if err := p.client.post(ctx, "/proofs/verify", body, &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// Get retrieves a proof by ID
func (p *ProofsClient) Get(ctx context.Context, proofID string) (*Proof, error) {
	var result Proof
	if err := p.client.get(ctx, "/proofs/"+proofID, &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// List lists recent proofs
func (p *ProofsClient) List(ctx context.Context, limit, offset int) ([]Proof, error) {
	path := "/proofs"
	if limit > 0 || offset > 0 {
		path = fmt.Sprintf("/proofs?limit=%d&offset=%d", limit, offset)
	}

	var result proofsResponse
	if err := p.client.get(ctx, path, &result); err != nil {
		return nil, err
	}
	return result.Proofs, nil
}

// Serialize serializes a proof to bytes
func (p *ProofsClient) Serialize(ctx context.Context, proof *Proof) ([]byte, error) {
	var result serializeResponse
	if err := p.client.post(ctx, "/proofs/serialize", map[string]string{"proof_id": proof.ID}, &result); err != nil {
		return nil, err
	}
	return base64.StdEncoding.DecodeString(result.Data)
}

// Deserialize deserializes bytes to a proof
func (p *ProofsClient) Deserialize(ctx context.Context, data []byte, system ProofSystem) (*Proof, error) {
	body := map[string]interface{}{
		"data":   base64.StdEncoding.EncodeToString(data),
		"system": system,
	}

	var result Proof
	if err := p.client.post(ctx, "/proofs/deserialize", body, &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// CircuitsClient handles circuit operations
type CircuitsClient struct {
	client *Client
}

// GetVerificationKey gets the verification key for a circuit
func (c *CircuitsClient) GetVerificationKey(ctx context.Context, circuitType CircuitType, system *ProofSystem) (*VerificationKey, error) {
	sys := c.client.DefaultProofSystem()
	if system != nil {
		sys = *system
	}

	var result VerificationKey
	if err := c.client.get(ctx, fmt.Sprintf("/circuits/%s/vk?system=%s", circuitType, sys), &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// GetProvingKey gets the proving key for a circuit
func (c *CircuitsClient) GetProvingKey(ctx context.Context, circuitType CircuitType, system *ProofSystem) (*ProvingKey, error) {
	sys := c.client.DefaultProofSystem()
	if system != nil {
		sys = *system
	}

	var result ProvingKey
	if err := c.client.get(ctx, fmt.Sprintf("/circuits/%s/pk?system=%s", circuitType, sys), &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// List lists available circuits
func (c *CircuitsClient) List(ctx context.Context) ([]CircuitConfig, error) {
	var result circuitsResponse
	if err := c.client.get(ctx, "/circuits", &result); err != nil {
		return nil, err
	}
	return result.Circuits, nil
}

// GetConfig gets circuit configuration
func (c *CircuitsClient) GetConfig(ctx context.Context, circuitType CircuitType) (*CircuitConfig, error) {
	var result CircuitConfig
	if err := c.client.get(ctx, fmt.Sprintf("/circuits/%s/config", circuitType), &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// Compile compiles a custom circuit
func (c *CircuitsClient) Compile(ctx context.Context, config CircuitConfig) (string, error) {
	var result circuitIDResponse
	if err := c.client.post(ctx, "/circuits/compile", config, &result); err != nil {
		return "", err
	}
	return result.CircuitID, nil
}

// RollupClient handles rollup operations
type RollupClient struct {
	client *Client
}

// GetStats gets rollup statistics
func (r *RollupClient) GetStats(ctx context.Context) (*RollupStats, error) {
	var result RollupStats
	if err := r.client.get(ctx, "/rollup/stats", &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// GetCurrentBatch gets the current batch
func (r *RollupClient) GetCurrentBatch(ctx context.Context) (*Batch, error) {
	var result Batch
	if err := r.client.get(ctx, "/rollup/batch/current", &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// GetBatch gets a batch by number
func (r *RollupClient) GetBatch(ctx context.Context, batchNumber uint64) (*Batch, error) {
	var result Batch
	if err := r.client.get(ctx, fmt.Sprintf("/rollup/batch/%d", batchNumber), &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// SubmitTransfer submits a transfer
func (r *RollupClient) SubmitTransfer(ctx context.Context, transfer Transfer) (string, error) {
	var result txIDResponse
	if err := r.client.post(ctx, "/rollup/transfer", transfer, &result); err != nil {
		return "", err
	}
	return result.TxID, nil
}

// SubmitDeposit submits a deposit
func (r *RollupClient) SubmitDeposit(ctx context.Context, deposit Deposit) (string, error) {
	var result txIDResponse
	if err := r.client.post(ctx, "/rollup/deposit", deposit, &result); err != nil {
		return "", err
	}
	return result.TxID, nil
}

// SubmitWithdrawal submits a withdrawal
func (r *RollupClient) SubmitWithdrawal(ctx context.Context, withdrawal Withdrawal) (string, error) {
	var result txIDResponse
	if err := r.client.post(ctx, "/rollup/withdraw", withdrawal, &result); err != nil {
		return "", err
	}
	return result.TxID, nil
}

// FinalizeBatch finalizes the current batch
func (r *RollupClient) FinalizeBatch(ctx context.Context) (*Batch, error) {
	var result Batch
	if err := r.client.post(ctx, "/rollup/batch/finalize", struct{}{}, &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// GetPendingTransactions gets pending transactions
func (r *RollupClient) GetPendingTransactions(ctx context.Context) ([]Transfer, error) {
	var result transactionsResponse
	if err := r.client.get(ctx, "/rollup/pending", &result); err != nil {
		return nil, err
	}
	return result.Transactions, nil
}

// StateClient handles state operations
type StateClient struct {
	client *Client
}

// GetRoot gets the current state root
func (s *StateClient) GetRoot(ctx context.Context) (string, error) {
	var result rootResponse
	if err := s.client.get(ctx, "/state/root", &result); err != nil {
		return "", err
	}
	return result.Root, nil
}

// GetAccount gets account state
func (s *StateClient) GetAccount(ctx context.Context, address string) (*AccountState, error) {
	var result AccountState
	if err := s.client.get(ctx, "/state/account/"+address, &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// GetMerkleProof gets merkle proof for account inclusion
func (s *StateClient) GetMerkleProof(ctx context.Context, address string) (*MerkleProof, error) {
	var result MerkleProof
	if err := s.client.get(ctx, "/state/proof/"+address, &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// VerifyMerkleProof verifies a merkle proof
func (s *StateClient) VerifyMerkleProof(ctx context.Context, proof MerkleProof) (bool, error) {
	var result validResponse
	if err := s.client.post(ctx, "/state/proof/verify", proof, &result); err != nil {
		return false, err
	}
	return result.Valid, nil
}

// GetStateAtBatch gets state at specific batch
func (s *StateClient) GetStateAtBatch(ctx context.Context, batchNumber uint64) (map[string]interface{}, error) {
	var result map[string]interface{}
	if err := s.client.get(ctx, fmt.Sprintf("/state/batch/%d", batchNumber), &result); err != nil {
		return nil, err
	}
	return result, nil
}

// CeremonyClient handles ceremony operations
type CeremonyClient struct {
	client *Client
}

// GetStatus gets ceremony status
func (c *CeremonyClient) GetStatus(ctx context.Context, circuitType CircuitType, system *ProofSystem) (*TrustedSetup, error) {
	sys := c.client.DefaultProofSystem()
	if system != nil {
		sys = *system
	}

	var result TrustedSetup
	if err := c.client.get(ctx, fmt.Sprintf("/ceremony/%s?system=%s", circuitType, sys), &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// Contribute contributes to ceremony
func (c *CeremonyClient) Contribute(ctx context.Context, circuitType CircuitType, entropy []byte, system *ProofSystem) (*CeremonyContribution, error) {
	sys := c.client.DefaultProofSystem()
	if system != nil {
		sys = *system
	}

	body := map[string]interface{}{
		"circuit_type": circuitType,
		"entropy":      base64.StdEncoding.EncodeToString(entropy),
		"system":       sys,
	}

	var result CeremonyContribution
	if err := c.client.post(ctx, "/ceremony/contribute", body, &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// VerifyContribution verifies a contribution
func (c *CeremonyClient) VerifyContribution(ctx context.Context, circuitType CircuitType, contributionHash string) (bool, error) {
	body := map[string]interface{}{
		"circuit_type":      circuitType,
		"contribution_hash": contributionHash,
	}

	var result validResponse
	if err := c.client.post(ctx, "/ceremony/verify", body, &result); err != nil {
		return false, err
	}
	return result.Valid, nil
}

// ListContributions lists contributions
func (c *CeremonyClient) ListContributions(ctx context.Context, circuitType CircuitType) ([]CeremonyContribution, error) {
	var result contributionsResponse
	if err := c.client.get(ctx, fmt.Sprintf("/ceremony/%s/contributions", circuitType), &result); err != nil {
		return nil, err
	}
	return result.Contributions, nil
}
