/*
Copyright 2024.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package providers

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"

	o11yv1 "github.com/tidbcloud/observability/api/regionalserver/resourceworker/v1"
)

// AKSCloudProvider implements CloudProviderInterface for Azure AKS via Azure Resource Manager REST API.
//
// Cluster identifier supports either:
//   - Full managed cluster resource ID:
//     "/subscriptions/<sub>/resourceGroups/<rg>/providers/Microsoft.ContainerService/managedClusters/<name>"
//   - Short cluster name with env vars AZURE_SUBSCRIPTION_ID and AZURE_RESOURCE_GROUP set.
type AKSCloudProvider struct {
	cred       azcore.TokenCredential
	httpClient *http.Client
	apiVersion string
	region     string
}

func NewAKSCloudProvider(region string) (*AKSCloudProvider, error) {
	cred, err := azidentity.NewDefaultAzureCredential(nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create Azure credential: %w", err)
	}
	return &AKSCloudProvider{
		cred:       cred,
		httpClient: &http.Client{Timeout: 60 * time.Second},
		apiVersion: "2024-05-01",
		region:     region,
	}, nil
}

func (a *AKSCloudProvider) DescribeCluster(ctx context.Context, clusterName string) (any, error) {
	clusterID, err := a.clusterResourceID(clusterName)
	if err != nil {
		return nil, err
	}
	var out any
	if err := a.armDoJSON(ctx, http.MethodGet, a.armURL(clusterID), nil, &out); err != nil {
		return nil, fmt.Errorf("failed to describe AKS cluster: %w", err)
	}
	return out, nil
}

func (a *AKSCloudProvider) UpgradeControlPlane(ctx context.Context, clusterName, targetVersion string) error {
	clusterID, err := a.clusterResourceID(clusterName)
	if err != nil {
		return err
	}
	body := map[string]any{
		"location": a.region,
		"properties": map[string]any{
			"kubernetesVersion": targetVersion,
		},
	}
	if err := a.armDoJSON(ctx, http.MethodPut, a.armURL(clusterID), body, nil); err != nil {
		return fmt.Errorf("failed to upgrade AKS control plane: %w", err)
	}
	return nil
}

func (_ *AKSCloudProvider) UpgradeClusterAddons(_ context.Context, _ string, _ []o11yv1.AddonUpgrade) error {
	return nil
}

func (_ *AKSCloudProvider) GetClusterAddonsStatus(_ context.Context, _ string, addons []o11yv1.AddonUpgrade) ([]AddonStatusInfo, error) {
	out := make([]AddonStatusInfo, len(addons))
	for i, a := range addons {
		out[i] = AddonStatusInfo{Name: a.Name, Version: a.Version, Status: "ACTIVE"}
	}
	return out, nil
}

func (_ *AKSCloudProvider) ListClusterAddons(_ context.Context, _ string) ([]AddonStatusInfo, error) {
	// AKS addons are not modeled the same way as EKS; return empty.
	return nil, nil
}

func (a *AKSCloudProvider) GetNodeGroupsByLabels(ctx context.Context, clusterName string, labelSelector map[string]string) ([]NodeGroup, error) {
	clusterID, err := a.clusterResourceID(clusterName)
	if err != nil {
		return nil, err
	}
	url := a.armURL(clusterID + "/agentPools")
	var resp struct {
		Value []aksAgentPool `json:"value"`
	}
	if err := a.armDoJSON(ctx, http.MethodGet, url, nil, &resp); err != nil {
		return nil, fmt.Errorf("failed to list AKS agentPools: %w", err)
	}

	var out []NodeGroup
	for _, ap := range resp.Value {
		ng := ap.toNodeGroup()
		matches := true
		for k, v := range labelSelector {
			if ng.Labels[k] != v {
				matches = false
				break
			}
		}
		if matches {
			out = append(out, ng)
		}
	}
	return out, nil
}

func (a *AKSCloudProvider) GetNodeGroupByName(ctx context.Context, clusterId string, nodeGroupName string) (*NodeGroup, error) {
	clusterID, err := a.clusterResourceID(clusterId)
	if err != nil {
		return nil, err
	}
	url := a.armURL(clusterID + "/agentPools/" + nodeGroupName)
	var ap aksAgentPool
	if err := a.armDoJSON(ctx, http.MethodGet, url, nil, &ap); err != nil {
		if strings.Contains(err.Error(), "Could not find the agentpool") {
			return nil, ErrNodeGroupNotFound
		}
		return nil, fmt.Errorf("failed to describe AKS agentPool %q: %w", nodeGroupName, err)
	}
	ng := ap.toNodeGroup()
	return &ng, nil
}

func (a *AKSCloudProvider) CreateNodeGroup(ctx context.Context, clusterName string, spec NodeGroup) (string, error) {
	clusterID, err := a.clusterResourceID(clusterName)
	if err != nil {
		return "", err
	}
	url := a.armURL(clusterID + "/agentPools/" + spec.Name)
	body := map[string]any{
		"properties": map[string]any{
			"count":               spec.DesiredSize,
			"vmSize":              spec.InstanceType[0],
			"mode":                "User",
			"osType":              "Linux",
			"orchestratorVersion": spec.Version,
			"nodeLabels":          spec.Labels,
		},
	}
	if spec.MinSize >= 0 && spec.MaxSize > 0 {
		props := body["properties"].(map[string]any)
		props["enableAutoScaling"] = true
		props["minCount"] = spec.MinSize
		props["maxCount"] = spec.MaxSize
	}
	if err := a.armDoJSON(ctx, http.MethodPut, url, body, nil); err != nil {
		return "", fmt.Errorf("failed to create AKS agentPool: %w", err)
	}
	return spec.Name, nil
}

func (a *AKSCloudProvider) UpdateNodeGroupMinSize(ctx context.Context, clusterName, nodeGroupID string, minSize int) error {
	ng, err := a.GetNodeGroupByName(ctx, clusterName, nodeGroupID)
	if err != nil {
		return err
	}
	clusterID, err := a.clusterResourceID(clusterName)
	if err != nil {
		return err
	}
	maxCount := ng.MaxSize
	if ng.MaxSize == 0 {
		maxCount = 99
	}
	body := map[string]any{
		"properties": map[string]any{
			"enableAutoScaling": true,
			"minCount":          minSize,
			"maxCount":          maxCount,
		},
	}
	if err := a.armDoJSON(ctx, http.MethodPatch, a.armURL(clusterID+"/agentPools/"+nodeGroupID), body, nil); err != nil {
		return fmt.Errorf("failed to update AKS agentPool %q min size: %w", nodeGroupID, err)
	}
	return nil
}

func (a *AKSCloudProvider) UpgradeNodeGroup(ctx context.Context, clusterName, nodeGroupID, targetVersion string, _ *o11yv1.NodeGroupUpgradeOptions) (string, error) {
	clusterID, err := a.clusterResourceID(clusterName)
	if err != nil {
		return "", err
	}
	url := a.armURL(clusterID + "/agentPools/" + nodeGroupID)
	body := map[string]any{
		"properties": map[string]any{
			"orchestratorVersion": targetVersion,
		},
	}
	if err := a.armDoJSON(ctx, http.MethodPut, url, body, nil); err != nil {
		return "", fmt.Errorf("failed to upgrade AKS agentPool %q: %w", nodeGroupID, err)
	}
	return "", nil
}

func (_ *AKSCloudProvider) GetNodeGroupUpgradeOperationStatus(_ context.Context, _, _, _ string) (bool, error) {
	return true, nil
}

func (a *AKSCloudProvider) DeleteNodeGroup(ctx context.Context, clusterName, nodeGroupID string) error {
	clusterID, err := a.clusterResourceID(clusterName)
	if err != nil {
		return err
	}
	url := a.armURL(clusterID + "/agentPools/" + nodeGroupID)
	if err := a.armDoJSON(ctx, http.MethodDelete, url, nil, nil); err != nil {
		return fmt.Errorf("failed to delete AKS agentPool %q: %w", nodeGroupID, err)
	}
	return nil
}

func (a *AKSCloudProvider) GetNodeGroupStatus(ctx context.Context, clusterName, nodeGroupID string) (NodeGroupStatus, error) {
	ng, err := a.GetNodeGroupByName(ctx, clusterName, nodeGroupID)
	if err != nil {
		return "", err
	}
	return ng.Status, nil
}

func (a *AKSCloudProvider) GetNodeGroupLabelSelector(ctx context.Context, clusterName, nodeGroupID string) (map[string]string, error) {
	_ = ctx
	_ = clusterName
	// AKS labels nodes with their agentpool name under kubernetes.azure.com/agentpool.
	return map[string]string{
		"kubernetes.azure.com/agentpool": nodeGroupID,
	}, nil
}

func (a *AKSCloudProvider) GetControlPlaneVersion(ctx context.Context, clusterName string) (string, error) {
	clusterID, err := a.clusterResourceID(clusterName)
	if err != nil {
		return "", err
	}
	var resp struct {
		Properties struct {
			KubernetesVersion string `json:"kubernetesVersion"`
		} `json:"properties"`
	}
	if err := a.armDoJSON(ctx, http.MethodGet, a.armURL(clusterID), nil, &resp); err != nil {
		return "", fmt.Errorf("failed to get AKS control plane version: %w", err)
	}
	if resp.Properties.KubernetesVersion == "" {
		return "", fmt.Errorf("AKS cluster %q kubernetesVersion is empty", clusterName)
	}
	return resp.Properties.KubernetesVersion, nil
}

// GetControlPlaneStatus returns the AKS cluster control plane status (ACTIVE or UPDATING).
func (a *AKSCloudProvider) GetControlPlaneStatus(ctx context.Context, clusterName string) (ControlPlaneStatus, error) {
	clusterID, err := a.clusterResourceID(clusterName)
	if err != nil {
		return "", err
	}
	var resp struct {
		Properties struct {
			ProvisioningState string `json:"provisioningState"`
		} `json:"properties"`
	}
	if err := a.armDoJSON(ctx, http.MethodGet, a.armURL(clusterID), nil, &resp); err != nil {
		return "", fmt.Errorf("failed to describe AKS cluster: %w", err)
	}
	// Updating = long-running update (e.g. version upgrade) in progress.
	if strings.EqualFold(resp.Properties.ProvisioningState, "Upgrading") {
		return ControlPlaneStatusUpdating, nil
	}
	if strings.EqualFold(resp.Properties.ProvisioningState, "Succeeded") {
		return ControlPlaneStatusActive, nil
	}
	return ControlPlaneStatusOther, nil
}

// GetInstanceTypeByTemplateVersion is not applicable to AKS (no EC2-style launch templates). Returns empty string, nil.
func (_ *AKSCloudProvider) GetInstanceTypeByTemplateVersion(_ context.Context, _, _, _ string) (string, error) {
	return "", nil
}

func (a *AKSCloudProvider) armURL(resourceID string) string {
	return fmt.Sprintf("https://management.azure.com%s?api-version=%s", resourceID, a.apiVersion)
}

func (a *AKSCloudProvider) clusterResourceID(clusterName string) (string, error) {
	if strings.HasPrefix(clusterName, "/subscriptions/") {
		return clusterName, nil
	}
	sub := os.Getenv("AZURE_SUBSCRIPTION_ID")
	rg := os.Getenv("AZURE_RESOURCE_GROUP")
	if sub == "" || rg == "" {
		return "", fmt.Errorf("AKS cluster name %q is not a resource ID; set env AZURE_SUBSCRIPTION_ID and AZURE_RESOURCE_GROUP or pass full resource ID", clusterName)
	}
	return fmt.Sprintf("/subscriptions/%s/resourceGroups/%s/providers/Microsoft.ContainerService/managedClusters/%s", sub, rg, clusterName), nil
}

func (a *AKSCloudProvider) armDoJSON(ctx context.Context, method, url string, in any, out any) error {
	var body io.Reader
	if in != nil {
		b, err := json.Marshal(in)
		if err != nil {
			return err
		}
		body = bytes.NewReader(b)
	}

	req, err := http.NewRequestWithContext(ctx, method, url, body)
	if err != nil {
		return err
	}
	req.Header.Set("Accept", "application/json")
	if in != nil {
		req.Header.Set("Content-Type", "application/json")
	}

	tok, err := a.cred.GetToken(ctx, policy.TokenRequestOptions{
		Scopes: []string{"https://management.azure.com/.default"},
	})
	if err != nil {
		return fmt.Errorf("failed to get Azure token: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+tok.Token)

	resp, err := a.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	// Azure ARM returns 202 for async operations; treat as success.
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		b, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("arm request failed: %s: %s", resp.Status, string(b))
	}

	if out == nil {
		return nil
	}
	b, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	if len(b) == 0 {
		return nil
	}
	return json.Unmarshal(b, out)
}

type aksAgentPool struct {
	Name       string `json:"name"`
	Properties struct {
		Count               int               `json:"count"`
		VMSize              string            `json:"vmSize"`
		OrchestratorVersion string            `json:"orchestratorVersion"`
		NodeLabels          map[string]string `json:"nodeLabels"`
		EnableAutoScaling   *bool             `json:"enableAutoScaling"`
		MinCount            *int              `json:"minCount"`
		MaxCount            *int              `json:"maxCount"`
		ProvisioningState   string            `json:"provisioningState"`
	} `json:"properties"`
}

func (ap *aksAgentPool) toNodeGroup() NodeGroup {
	ng := NodeGroup{
		ID:           ap.Name,
		Name:         ap.Name,
		Labels:       map[string]string{},
		Version:      ap.Properties.OrchestratorVersion,
		InstanceType: []string{ap.Properties.VMSize},
		NodeCount:    ap.Properties.Count,
		DesiredSize:  ap.Properties.Count,
	}
	for k, v := range ap.Properties.NodeLabels {
		ng.Labels[k] = v
	}
	if ap.Properties.EnableAutoScaling != nil && *ap.Properties.EnableAutoScaling {
		if ap.Properties.MinCount != nil {
			ng.MinSize = *ap.Properties.MinCount
		}
		if ap.Properties.MaxCount != nil {
			ng.MaxSize = *ap.Properties.MaxCount
		}
	}

	switch strings.ToLower(ap.Properties.ProvisioningState) {
	case "succeeded":
		ng.Status = NodeGroupStatusActive
	case "upgrading":
		ng.Status = NodeGroupStatusUpdating
	case "creating":
		ng.Status = NodeGroupStatusCreating
	case "deleting":
		ng.Status = NodeGroupStatusDeleting
	case "failed":
		ng.Status = NodeGroupStatusCreateFailed
	default:
		ng.Status = NodeGroupStatus(ap.Properties.ProvisioningState)
	}

	return ng
}