// Copyright 2025 Microsoft Corporation
//
// 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 metrics

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/go-logr/logr"
	"github.com/prometheus/client_golang/prometheus"

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

	"github.com/Azure/ARO-HCP/tooling/aro-hcp-exporter/internal/cluster"
	"github.com/Azure/ARO-HCP/tooling/aro-hcp-exporter/pkg/graphquery"
	"github.com/Azure/ARO-HCP/tooling/metricscache"
)

const (
	ServiceTagUsageCollectorName = "service-tag-usage"
)

var (
	// ServiceTagUsageByPublicIpCountDesc is the descriptor for the public IP count metric by service tag
	ServiceTagUsageByPublicIpCountDesc = prometheus.NewDesc(
		"public_ip_count_by_region_service_tag",
		"Number of public IP addresses in a region by service tag",
		[]string{"subscription_id", "region", "service_tag_type", "service_tag_value"},
		nil,
	)
)

// ServiceTagUsageCollector is a Prometheus collector that gathers public IP metrics from Azure
type ServiceTagUsageCollector struct {
	clusterClient *cluster.ClusterDiscoveryPoller
	credential    azcore.TokenCredential
	cache         *metricscache.Cache
	region        string
	errorCounter  prometheus.Counter
}

var _ CachingCollector = &ServiceTagUsageCollector{}

// NewServiceTagUsageCollector creates a new ServiceTagUsageCollector
func NewServiceTagUsageCollector(clusterClient *cluster.ClusterDiscoveryPoller, region string, credential azcore.TokenCredential, cacheTTL time.Duration, errorCounter prometheus.Counter) *ServiceTagUsageCollector {
	return &ServiceTagUsageCollector{
		clusterClient: clusterClient,
		credential:    credential,
		cache:         metricscache.NewCache(cacheTTL),
		region:        region,
		errorCounter:  errorCounter,
	}
}

func (c *ServiceTagUsageCollector) Name() string {
	return ServiceTagUsageCollectorName
}

func (c *ServiceTagUsageCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- ServiceTagUsageByPublicIpCountDesc
}

func (c *ServiceTagUsageCollector) Collect(ch chan<- prometheus.Metric) {
	for _, metric := range c.cache.GetAll() {
		ch <- metric
	}
}

type IPTag struct {
	ServiceTagType  string `json:"ipTagType"`
	ServiceTagValue string `json:"tag"`
}

type PublicIPAddress struct {
	Location       string
	IpTagsString   string
	SubscriptionId string
	Count          float64
}

func buildIPQuery(region string) string {
	return fmt.Sprintf(`
resources
| where type =~ 'microsoft.network/publicipaddresses'
| where location =~ '%s'
| extend ipTagsString = tostring(properties['ipTags'])
| summarize Count=count() by subscriptionId, location, ipTagsString
`, graphquery.EscapeKQL(region))
}

func parseIPTags(ipTagsAsString string) ([]IPTag, error) {
	ipTags := []IPTag{}

	if ipTagsAsString == "" {
		return ipTags, nil
	}
	err := json.Unmarshal([]byte(ipTagsAsString), &ipTags)
	if err != nil {
		return nil, fmt.Errorf("error parsing IPs %s, %w", ipTagsAsString, err)
	}
	return ipTags, nil
}

func (c *ServiceTagUsageCollector) CollectMetricValues(ctx context.Context) {
	logger := logr.FromContextOrDiscard(ctx)

	discoverResult := c.clusterClient.GetDiscoverResult(ctx)
	subscriptionIDsPtrs := make([]*string, 0, len(discoverResult.SubscriptionIDs))
	for _, id := range discoverResult.SubscriptionIDs {
		subscriptionIDsPtrs = append(subscriptionIDsPtrs, to.Ptr(id))
	}
	if len(subscriptionIDsPtrs) == 0 {
		logger.Info("No subscriptions discovered, skipping service tag usage collection")
		return
	}
	client, err := graphquery.NewResourceGraphClient(c.credential, subscriptionIDsPtrs)
	if err != nil {
		c.errorCounter.Inc()
		logger.Error(err, "failed to create Resource Graph client")
		return
	}

	var publicIPs []PublicIPAddress

	q := buildIPQuery(c.region)
	err = client.ExecuteConvertRequest(ctx, graphquery.ResourceGraphRequest{
		Query:  &q,
		Output: &publicIPs,
	})
	if err != nil {
		c.errorCounter.Inc()
		logger.Error(err, "error collecting public IP addresses")
		return
	}

	for _, publicIP := range publicIPs {
		var ipTags []IPTag
		if publicIP.IpTagsString != "[]" {
			ipTags, err = parseIPTags(publicIP.IpTagsString)
			if err != nil {
				c.errorCounter.Inc()
				logger.Error(err, "error parsing IP tags")
				continue
			}
		} else {
			ipTags = []IPTag{
				{
					ServiceTagType:  "not-set",
					ServiceTagValue: "not-set",
				},
			}
		}
		for _, ipTag := range ipTags {
			labels := []string{publicIP.SubscriptionId, publicIP.Location, ipTag.ServiceTagType, ipTag.ServiceTagValue}
			c.cache.Set(strings.Join(labels, "/"), prometheus.MustNewConstMetric(
				ServiceTagUsageByPublicIpCountDesc,
				prometheus.GaugeValue,
				publicIP.Count,
				labels...,
			))
		}
	}
}
