package capacitycontroller

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"
)

type KubernetesConfig struct {
	Namespace  string
	APIServer  string
	TokenFile  string
	CAFile     string
	HTTPClient *http.Client
}

type KubernetesClient struct {
	namespace string
	baseURL   string
	tokenFile string
	client    *http.Client
}

type KubernetesSiteState struct {
	Site                     string
	NodeName                 string
	CurrentReplicas          int
	Revision                 string
	Image                    string
	AllocatableMemoryBytes   int64
	UsedMemoryBytes          int64
	WorkerRequestMemoryBytes int64
	AllocatableMilliCPU      int64
	UsedMilliCPU             int64
	WorkerRequestMilliCPU    int64
	MemoryPressure           bool
	Nodes                    []KubernetesNodeState
}

type KubernetesNodeState struct {
	NodeName               string
	CurrentReplicas        int
	AllocatableMemoryBytes int64
	UsedMemoryBytes        int64
	WorkerUsedMemoryBytes  int64
	AllocatableMilliCPU    int64
	UsedMilliCPU           int64
	WorkerUsedMilliCPU     int64
	MemoryPressure         bool
}

func NewKubernetesClient(config KubernetesConfig) (*KubernetesClient, error) {
	if strings.TrimSpace(config.Namespace) == "" {
		return nil, errors.New("capacity controller Kubernetes namespace is required")
	}
	if config.APIServer == "" {
		config.APIServer = "https://kubernetes.default.svc"
	}
	if config.TokenFile == "" {
		config.TokenFile = "/var/run/secrets/kubernetes.io/serviceaccount/token"
	}
	parsed, err := url.Parse(strings.TrimRight(config.APIServer, "/"))
	if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" ||
		(parsed.Scheme != "https" && config.HTTPClient == nil) {
		return nil, errors.New("capacity controller Kubernetes API server URL is invalid")
	}
	client := config.HTTPClient
	if client == nil {
		if config.CAFile == "" {
			config.CAFile = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
		}
		certificate, readErr := os.ReadFile(config.CAFile)
		if readErr != nil {
			return nil, fmt.Errorf("read Kubernetes service account CA: %w", readErr)
		}
		roots := x509.NewCertPool()
		if !roots.AppendCertsFromPEM(certificate) {
			return nil, errors.New("Kubernetes service account CA is invalid")
		}
		transport := http.DefaultTransport.(*http.Transport).Clone()
		transport.Proxy = nil
		transport.DisableCompression = true
		transport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12, RootCAs: roots}
		client = &http.Client{
			Timeout:   10 * time.Second,
			Transport: transport,
			CheckRedirect: func(*http.Request, []*http.Request) error {
				return http.ErrUseLastResponse
			},
		}
	}
	return &KubernetesClient{
		namespace: config.Namespace,
		baseURL:   strings.TrimRight(config.APIServer, "/"),
		tokenFile: config.TokenFile,
		client:    client,
	}, nil
}

func (client *KubernetesClient) SiteState(ctx context.Context, site string) (KubernetesSiteState, error) {
	site = strings.TrimSpace(site)
	if site == "" {
		return KubernetesSiteState{}, errors.New("site is required")
	}
	var deployment struct {
		Metadata struct {
			Name string `json:"name"`
		} `json:"metadata"`
		Spec struct {
			Replicas int `json:"replicas"`
			Template struct {
				Spec struct {
					Containers []struct {
						Name  string `json:"name"`
						Image string `json:"image"`
						Env   []struct {
							Name  string `json:"name"`
							Value string `json:"value"`
						} `json:"env"`
						Resources struct {
							Requests map[string]string `json:"requests"`
						} `json:"resources"`
					} `json:"containers"`
				} `json:"spec"`
			} `json:"template"`
		} `json:"spec"`
	}
	deploymentPath := fmt.Sprintf(
		"/apis/apps/v1/namespaces/%s/deployments/easyai-worker-%s",
		url.PathEscape(client.namespace),
		url.PathEscape(site),
	)
	if err := client.getJSON(ctx, deploymentPath, &deployment); err != nil {
		return KubernetesSiteState{}, err
	}
	state := KubernetesSiteState{Site: site, CurrentReplicas: deployment.Spec.Replicas}
	for _, container := range deployment.Spec.Template.Spec.Containers {
		if container.Name != "worker" {
			continue
		}
		state.Image = container.Image
		state.WorkerRequestMemoryBytes, _ = parseBinaryQuantity(container.Resources.Requests["memory"])
		state.WorkerRequestMilliCPU, _ = parseCPUQuantity(container.Resources.Requests["cpu"])
		for _, environment := range container.Env {
			if environment.Name == "AI_GATEWAY_REVISION" {
				state.Revision = environment.Value
			}
		}
	}
	var nodes struct {
		Items []struct {
			Metadata struct {
				Name string `json:"name"`
			} `json:"metadata"`
			Status struct {
				Allocatable map[string]string `json:"allocatable"`
				Conditions  []struct {
					Type   string `json:"type"`
					Status string `json:"status"`
				} `json:"conditions"`
			} `json:"status"`
		} `json:"items"`
	}
	nodePath := "/api/v1/nodes?labelSelector=" + url.QueryEscape(
		"easyai.io/site="+site+",easyai.io/worker=true",
	)
	if err := client.getJSON(ctx, nodePath, &nodes); err != nil {
		return KubernetesSiteState{}, err
	}
	if len(nodes.Items) == 0 {
		return KubernetesSiteState{}, fmt.Errorf("site %s has no matching nodes", site)
	}
	var pods struct {
		Items []struct {
			Metadata struct {
				Name              string     `json:"name"`
				DeletionTimestamp *time.Time `json:"deletionTimestamp"`
			} `json:"metadata"`
			Spec struct {
				NodeName string `json:"nodeName"`
			} `json:"spec"`
		} `json:"items"`
	}
	podPath := fmt.Sprintf(
		"/api/v1/namespaces/%s/pods?labelSelector=%s",
		url.PathEscape(client.namespace),
		url.QueryEscape("app.kubernetes.io/name=easyai-worker,easyai.io/site="+site),
	)
	if err := client.getJSON(ctx, podPath, &pods); err != nil {
		return KubernetesSiteState{}, err
	}
	replicasByNode := make(map[string]int, len(nodes.Items))
	nodeByPod := make(map[string]string, len(pods.Items))
	for _, pod := range pods.Items {
		if pod.Metadata.DeletionTimestamp == nil && pod.Spec.NodeName != "" {
			replicasByNode[pod.Spec.NodeName]++
			nodeByPod[pod.Metadata.Name] = pod.Spec.NodeName
		}
	}
	var podMetrics struct {
		Items []struct {
			Metadata struct {
				Name string `json:"name"`
			} `json:"metadata"`
			Containers []struct {
				Usage map[string]string `json:"usage"`
			} `json:"containers"`
		} `json:"items"`
	}
	podMetricsPath := fmt.Sprintf(
		"/apis/metrics.k8s.io/v1beta1/namespaces/%s/pods?labelSelector=%s",
		url.PathEscape(client.namespace),
		url.QueryEscape("app.kubernetes.io/name=easyai-worker,easyai.io/site="+site),
	)
	if err := client.getJSON(ctx, podMetricsPath, &podMetrics); err != nil {
		return KubernetesSiteState{}, err
	}
	workerMemoryByNode := make(map[string]int64, len(nodes.Items))
	workerCPUByNode := make(map[string]int64, len(nodes.Items))
	measuredPods := make(map[string]bool, len(podMetrics.Items))
	for _, pod := range podMetrics.Items {
		nodeName := nodeByPod[pod.Metadata.Name]
		if nodeName == "" {
			continue
		}
		measuredPods[pod.Metadata.Name] = true
		for _, container := range pod.Containers {
			memory, memoryErr := parseBinaryQuantity(container.Usage["memory"])
			cpu, cpuErr := parseCPUQuantity(container.Usage["cpu"])
			if memoryErr != nil || cpuErr != nil {
				return KubernetesSiteState{}, fmt.Errorf(
					"site %s pod %s returned incomplete resource metrics",
					site,
					pod.Metadata.Name,
				)
			}
			workerMemoryByNode[nodeName] += memory
			workerCPUByNode[nodeName] += cpu
		}
	}
	for podName := range nodeByPod {
		if !measuredPods[podName] {
			return KubernetesSiteState{}, fmt.Errorf(
				"site %s pod %s has no resource metrics yet",
				site,
				podName,
			)
		}
	}
	for _, node := range nodes.Items {
		nodeState := KubernetesNodeState{
			NodeName:              node.Metadata.Name,
			CurrentReplicas:       replicasByNode[node.Metadata.Name],
			WorkerUsedMemoryBytes: workerMemoryByNode[node.Metadata.Name],
			WorkerUsedMilliCPU:    workerCPUByNode[node.Metadata.Name],
		}
		nodeState.AllocatableMemoryBytes, _ = parseBinaryQuantity(node.Status.Allocatable["memory"])
		nodeState.AllocatableMilliCPU, _ = parseCPUQuantity(node.Status.Allocatable["cpu"])
		for _, condition := range node.Status.Conditions {
			if condition.Type == "MemoryPressure" && condition.Status == "True" {
				nodeState.MemoryPressure = true
				state.MemoryPressure = true
			}
		}
		var metrics struct {
			Usage map[string]string `json:"usage"`
		}
		metricsPath := "/apis/metrics.k8s.io/v1beta1/nodes/" + url.PathEscape(nodeState.NodeName)
		if err := client.getJSON(ctx, metricsPath, &metrics); err != nil {
			return KubernetesSiteState{}, err
		}
		nodeState.UsedMemoryBytes, _ = parseBinaryQuantity(metrics.Usage["memory"])
		nodeState.UsedMilliCPU, _ = parseCPUQuantity(metrics.Usage["cpu"])
		if nodeState.AllocatableMemoryBytes <= 0 || nodeState.AllocatableMilliCPU <= 0 {
			return KubernetesSiteState{}, fmt.Errorf(
				"site %s node %s returned incomplete resource quantities",
				site,
				nodeState.NodeName,
			)
		}
		state.Nodes = append(state.Nodes, nodeState)
		state.AllocatableMemoryBytes += nodeState.AllocatableMemoryBytes
		state.AllocatableMilliCPU += nodeState.AllocatableMilliCPU
		state.UsedMemoryBytes += nodeState.UsedMemoryBytes
		state.UsedMilliCPU += nodeState.UsedMilliCPU
	}
	if len(state.Nodes) == 1 {
		state.NodeName = state.Nodes[0].NodeName
	}
	if state.WorkerRequestMemoryBytes <= 0 || state.WorkerRequestMilliCPU <= 0 ||
		state.AllocatableMemoryBytes <= 0 || state.AllocatableMilliCPU <= 0 {
		return KubernetesSiteState{}, fmt.Errorf("site %s returned incomplete resource quantities", site)
	}
	return state, nil
}

func (client *KubernetesClient) ScaleWorkerDeployment(ctx context.Context, site string, replicas int) error {
	if replicas < 0 || replicas > 64 {
		return errors.New("worker replica target must be between 0 and 64")
	}
	path := fmt.Sprintf(
		"/apis/apps/v1/namespaces/%s/deployments/easyai-worker-%s/scale",
		url.PathEscape(client.namespace),
		url.PathEscape(strings.TrimSpace(site)),
	)
	return client.patchJSON(ctx, path, map[string]any{"spec": map[string]any{"replicas": replicas}})
}

func (client *KubernetesClient) SetPodDeletionCost(ctx context.Context, podName string, cost int) error {
	podName = strings.TrimSpace(podName)
	if podName == "" {
		return errors.New("worker pod name is required")
	}
	path := fmt.Sprintf(
		"/api/v1/namespaces/%s/pods/%s",
		url.PathEscape(client.namespace),
		url.PathEscape(podName),
	)
	return client.patchJSON(ctx, path, map[string]any{
		"metadata": map[string]any{
			"annotations": map[string]any{
				"controller.kubernetes.io/pod-deletion-cost": strconv.Itoa(cost),
			},
		},
	})
}

func (client *KubernetesClient) getJSON(ctx context.Context, path string, target any) error {
	request, err := client.request(ctx, http.MethodGet, path, nil)
	if err != nil {
		return err
	}
	response, err := client.client.Do(request)
	if err != nil {
		return fmt.Errorf("request Kubernetes API: %w", err)
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
		return fmt.Errorf("Kubernetes API %s returned HTTP %d", path, response.StatusCode)
	}
	decoder := json.NewDecoder(io.LimitReader(response.Body, 4*1024*1024))
	if err := decoder.Decode(target); err != nil {
		return fmt.Errorf("decode Kubernetes API %s: %w", path, err)
	}
	return nil
}

func (client *KubernetesClient) patchJSON(ctx context.Context, path string, payload any) error {
	encoded, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	request, err := client.request(ctx, http.MethodPatch, path, bytes.NewReader(encoded))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/merge-patch+json")
	response, err := client.client.Do(request)
	if err != nil {
		return fmt.Errorf("patch Kubernetes API: %w", err)
	}
	defer response.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return fmt.Errorf("Kubernetes API %s returned HTTP %d", path, response.StatusCode)
	}
	return nil
}

func (client *KubernetesClient) request(ctx context.Context, method string, path string, body io.Reader) (*http.Request, error) {
	token, err := os.ReadFile(client.tokenFile)
	if err != nil || strings.TrimSpace(string(token)) == "" {
		clear(token)
		return nil, errors.New("Kubernetes service account token is unavailable")
	}
	request, err := http.NewRequestWithContext(ctx, method, client.baseURL+path, body)
	if err != nil {
		clear(token)
		return nil, err
	}
	request.Header.Set("Authorization", "Bearer "+strings.TrimSpace(string(token)))
	request.Header.Set("Accept", "application/json")
	clear(token)
	return request, nil
}

func parseBinaryQuantity(value string) (int64, error) {
	value = strings.TrimSpace(value)
	if value == "" {
		return 0, errors.New("quantity is empty")
	}
	units := []struct {
		suffix     string
		multiplier int64
	}{
		{"Ei", 1 << 60},
		{"Pi", 1 << 50},
		{"Ti", 1 << 40},
		{"Gi", 1 << 30},
		{"Mi", 1 << 20},
		{"Ki", 1 << 10},
		{"G", 1_000_000_000},
		{"M", 1_000_000},
		{"K", 1_000},
	}
	for _, unit := range units {
		if strings.HasSuffix(value, unit.suffix) {
			number, err := strconv.ParseFloat(strings.TrimSuffix(value, unit.suffix), 64)
			return int64(number * float64(unit.multiplier)), err
		}
	}
	return strconv.ParseInt(value, 10, 64)
}

func parseCPUQuantity(value string) (int64, error) {
	value = strings.TrimSpace(value)
	if strings.HasSuffix(value, "n") {
		nanocores, err := strconv.ParseInt(strings.TrimSuffix(value, "n"), 10, 64)
		return nanocores / 1_000_000, err
	}
	if strings.HasSuffix(value, "u") {
		microcores, err := strconv.ParseInt(strings.TrimSuffix(value, "u"), 10, 64)
		return microcores / 1_000, err
	}
	if strings.HasSuffix(value, "m") {
		return strconv.ParseInt(strings.TrimSuffix(value, "m"), 10, 64)
	}
	cores, err := strconv.ParseFloat(value, 64)
	return int64(cores * 1000), err
}