// Package synor provides a Go SDK for Synor Compute.
//
// Access distributed heterogeneous compute resources (CPU, GPU, TPU, NPU, LPU)
// for AI/ML workloads at 90% cost reduction compared to traditional cloud.
//
// Example:
//
//	client := synor.NewClient("your-api-key")
//	result, err := client.MatMul(ctx, a, b, synor.FP16)
package synor

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

// Version of the SDK.
const Version = "0.1.0"

// DefaultEndpoint is the default API endpoint.
const DefaultEndpoint = "https://compute.synor.cc/api/v1"

// ProcessorType represents supported processor types.
type ProcessorType string

const (
	CPU    ProcessorType = "cpu"
	GPU    ProcessorType = "gpu"
	TPU    ProcessorType = "tpu"
	NPU    ProcessorType = "npu"
	LPU    ProcessorType = "lpu"
	FPGA   ProcessorType = "fpga"
	WASM   ProcessorType = "wasm"
	WebGPU ProcessorType = "webgpu"
)

// Precision represents computation precision levels.
type Precision string

const (
	FP64 Precision = "fp64"
	FP32 Precision = "fp32"
	FP16 Precision = "fp16"
	BF16 Precision = "bf16"
	INT8 Precision = "int8"
	INT4 Precision = "int4"
)

// BalancingStrategy represents job scheduling strategies.
type BalancingStrategy string

const (
	Speed    BalancingStrategy = "speed"
	Cost     BalancingStrategy = "cost"
	Energy   BalancingStrategy = "energy"
	Latency  BalancingStrategy = "latency"
	Balanced BalancingStrategy = "balanced"
)

// JobStatus represents job states.
type JobStatus string

const (
	Pending   JobStatus = "pending"
	Queued    JobStatus = "queued"
	Running   JobStatus = "running"
	Completed JobStatus = "completed"
	Failed    JobStatus = "failed"
	Cancelled JobStatus = "cancelled"
)

// Config holds client configuration.
type Config struct {
	APIKey    string
	Endpoint  string
	Strategy  BalancingStrategy
	Precision Precision
	Timeout   time.Duration
	Debug     bool
}

// DefaultConfig returns a default configuration.
func DefaultConfig(apiKey string) Config {
	return Config{
		APIKey:    apiKey,
		Endpoint:  DefaultEndpoint,
		Strategy:  Balanced,
		Precision: FP32,
		Timeout:   30 * time.Second,
		Debug:     false,
	}
}

// Client is the main Synor Compute client.
type Client struct {
	config     Config
	httpClient *http.Client
}

// NewClient creates a new client with the given API key.
func NewClient(apiKey string) *Client {
	return NewClientWithConfig(DefaultConfig(apiKey))
}

// NewClientWithConfig creates a new client with custom configuration.
func NewClientWithConfig(config Config) *Client {
	return &Client{
		config: config,
		httpClient: &http.Client{
			Timeout: config.Timeout,
		},
	}
}

// JobResult represents the result of a job.
type JobResult struct {
	JobID   string      `json:"job_id"`
	Status  JobStatus   `json:"status"`
	Data    interface{} `json:"data,omitempty"`
	Error   string      `json:"error,omitempty"`
	Metrics *JobMetrics `json:"metrics,omitempty"`
}

// JobMetrics contains execution metrics.
type JobMetrics struct {
	ExecutionTimeMs float64       `json:"execution_time_ms"`
	QueueTimeMs     float64       `json:"queue_time_ms"`
	ProcessorType   ProcessorType `json:"processor_type"`
	ProcessorID     string        `json:"processor_id"`
	FLOPS           float64       `json:"flops"`
	MemoryBytes     int64         `json:"memory_bytes"`
	CostMicro       int64         `json:"cost_micro"`
	EnergyMJ        float64       `json:"energy_mj"`
}

// Tensor represents a multi-dimensional array.
type Tensor struct {
	Data  []float32 `json:"data"`
	Shape []int     `json:"shape"`
	DType Precision `json:"dtype"`
}

// NewTensor creates a new tensor from data and shape.
func NewTensor(data []float32, shape []int, dtype Precision) *Tensor {
	return &Tensor{
		Data:  data,
		Shape: shape,
		DType: dtype,
	}
}

// Zeros creates a tensor of zeros.
func Zeros(shape []int, dtype Precision) *Tensor {
	size := 1
	for _, dim := range shape {
		size *= dim
	}
	return &Tensor{
		Data:  make([]float32, size),
		Shape: shape,
		DType: dtype,
	}
}

// Serialize converts tensor to transmission format.
func (t *Tensor) Serialize() map[string]interface{} {
	buf := new(bytes.Buffer)
	for _, v := range t.Data {
		var b [4]byte
		*(*float32)((&b[0])) = v
		buf.Write(b[:])
	}
	return map[string]interface{}{
		"data":  base64.StdEncoding.EncodeToString(buf.Bytes()),
		"shape": t.Shape,
		"dtype": t.DType,
	}
}

// MatMul performs matrix multiplication.
func (c *Client) MatMul(ctx context.Context, a, b *Tensor, precision Precision) (*JobResult, error) {
	job, err := c.SubmitJob(ctx, "matmul", map[string]interface{}{
		"a":         a.Serialize(),
		"b":         b.Serialize(),
		"precision": precision,
	})
	if err != nil {
		return nil, err
	}
	return job.Wait(ctx)
}

// Inference runs model inference.
func (c *Client) Inference(ctx context.Context, model, input string, maxTokens int) (*JobResult, error) {
	job, err := c.SubmitJob(ctx, "inference", map[string]interface{}{
		"model":      model,
		"input":      input,
		"max_tokens": maxTokens,
		"strategy":   c.config.Strategy,
	})
	if err != nil {
		return nil, err
	}
	return job.Wait(ctx)
}

// Job represents a submitted compute job.
type Job struct {
	ID     string
	client *Client
	status JobStatus
}

// SubmitJob submits a new compute job.
func (c *Client) SubmitJob(ctx context.Context, operation string, params map[string]interface{}) (*Job, error) {
	body := map[string]interface{}{
		"operation": operation,
		"params":    params,
		"strategy":  c.config.Strategy,
	}

	var resp struct {
		JobID string `json:"job_id"`
	}

	if err := c.request(ctx, "POST", "/jobs", body, &resp); err != nil {
		return nil, err
	}

	return &Job{
		ID:     resp.JobID,
		client: c,
		status: Pending,
	}, nil
}

// Wait waits for the job to complete.
func (j *Job) Wait(ctx context.Context) (*JobResult, error) {
	ticker := time.NewTicker(500 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-ticker.C:
			result, err := j.client.GetJobStatus(ctx, j.ID)
			if err != nil {
				return nil, err
			}
			if result.Status == Completed || result.Status == Failed || result.Status == Cancelled {
				return result, nil
			}
		}
	}
}

// GetJobStatus gets the status of a job.
func (c *Client) GetJobStatus(ctx context.Context, jobID string) (*JobResult, error) {
	var result JobResult
	if err := c.request(ctx, "GET", "/jobs/"+jobID, nil, &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// CancelJob cancels a running job.
func (c *Client) CancelJob(ctx context.Context, jobID string) error {
	return c.request(ctx, "DELETE", "/jobs/"+jobID, nil, nil)
}

// PricingInfo contains pricing information.
type PricingInfo struct {
	ProcessorType  ProcessorType `json:"processor_type"`
	SpotPrice      float64       `json:"spot_price"`
	AvgPrice24h    float64       `json:"avg_price_24h"`
	AWSEquivalent  float64       `json:"aws_equivalent"`
	SavingsPercent float64       `json:"savings_percent"`
}

// GetPricing gets current pricing for all processor types.
func (c *Client) GetPricing(ctx context.Context) ([]PricingInfo, error) {
	var resp struct {
		Pricing []PricingInfo `json:"pricing"`
	}
	if err := c.request(ctx, "GET", "/pricing", nil, &resp); err != nil {
		return nil, err
	}
	return resp.Pricing, nil
}

func (c *Client) request(ctx context.Context, method, path string, body interface{}, result interface{}) error {
	url := c.config.Endpoint + path

	var bodyReader io.Reader
	if body != nil {
		bodyBytes, err := json.Marshal(body)
		if err != nil {
			return fmt.Errorf("failed to marshal request: %w", err)
		}
		bodyReader = bytes.NewReader(bodyBytes)
	}

	req, err := http.NewRequestWithContext(ctx, method, url, bodyReader)
	if err != nil {
		return fmt.Errorf("failed to create request: %w", err)
	}

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

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		var errResp struct {
			Message string `json:"message"`
		}
		json.NewDecoder(resp.Body).Decode(&errResp)
		return &SynorError{
			Message:    errResp.Message,
			StatusCode: resp.StatusCode,
		}
	}

	if result != nil {
		if err := json.NewDecoder(resp.Body).Decode(result); err != nil {
			return fmt.Errorf("failed to decode response: %w", err)
		}
	}

	return nil
}

// SynorError represents an API error.
type SynorError struct {
	Message    string
	StatusCode int
}

func (e *SynorError) Error() string {
	return fmt.Sprintf("synor: %s (status %d)", e.Message, e.StatusCode)
}