/*
 * Copyright 2021 PingCAP, Inc.
 *
 * 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,
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"context"
	"github.com/PagerDuty/go-pagerduty"
	"gopkg.in/yaml.v3"
	"log"
	"strings"
	"time"
)

const flinkTSFormat string = "2006-01-02 15:04:05"

// Alert is the message struct send to sqs.
type Alert struct {
	AlertID    string `json:"alert_id"`
	TenantID   string `json:"tenant_id"`
	ClusterID  string `json:"cluster_id"`
	Name       string `json:"name"`
	IncidentID string `json:"incident_id"`

	TriggeredAt        string `json:"triggered_at"`
	ResolvedAt         string `json:"resolved_at,omitempty"`
	AcknowledgedAt     string `json:"acknowledged_at,omitempty"`
	TriggerAtUnix      int64  `json:"triggered_at_unix"`
	ResolveAtUnix      int64  `json:"resolved_at_unix"`
	AcknowledgedAtUnix int64  `json:"acknowledged_at_unix"`

	Severity string `json:"severity"`
	Urgency  string `json:"urgency"`
	Priority string `json:"priority"`

	Service          string `json:"service"`
	Team             string `json:"team"`
	AssignTo         string `json:"assign_to"`
	EscalationPolicy string `json:"escalation_policy"`
}

type collector struct {
	sink       chan Alert
	st         time.Time
	et         time.Time
	drainSince string
	drainUntil string
	pdClient   *pagerduty.Client
}

func New(ak string, st time.Time, et time.Time) *collector {
	c := &collector{
		sink:     make(chan Alert, 300),
		st:       st,
		et:       et,
		pdClient: pagerduty.NewClient(ak),
	}
	return c
}

func (c *collector) SetDrainRange(since string, until string) {
	c.drainSince = since
	c.drainUntil = until
}

func (c *collector) GetSink() <-chan Alert {
	return c.sink
}

// collectAlert call pagerduty rest api and extract information from target fields.
func (c *collector) collectAlert(ctx context.Context, opts pagerduty.ListIncidentsOptions) error {
	var hasMore = true
	for hasMore {
		ret, err := c.pdClient.ListIncidentsWithContext(ctx, opts)
		if err != nil {
			return err
		}
		hasMore = ret.More
		if hasMore {
			// update the offset
			opts.Offset += uint(len(ret.Incidents))
		}
		for _, incident := range ret.Incidents {
			var acknowledgeUnix int64 = 0
			var acknowledgeUTC time.Time
			var service string
			var team string
			var assignTo string
			var ePolicy string
			var priority string
			// no alert belong to this incident, skip it directly.
			if incident.AlertCounts.All == 0 {
				continue
			}
			alerts, err := c.pdClient.ListIncidentAlertsWithContext(ctx, incident.Id, pagerduty.ListIncidentAlertsOptions{})
			if err != nil {
				log.Printf("ListIncidentAlerts fail with indicdent id %+v\n", incident.Id)
				return err
			}
			// retrieve assign_to and acknowledge from different palce
			switch incident.Status {
			case "resolved":
				// get log entry and parse acknowledged_at and assign_to
				data, err := c.pdClient.ListIncidentLogEntriesWithContext(ctx, incident.Id, pagerduty.ListIncidentLogEntriesOptions{
				IsOverview:    true,
				})
				if err != nil {
					log.Printf("ListIncidentLogEntries fail with indicdent id %+v\n", incident.Id)
					return err
				}
				for _, ent := range data.LogEntries {
					switch ent.Type {
					case "resolve_log_entry":
						if ent.Agent.Type == "user_reference" || ent.Agent.Type == "service_reference" {
							assignTo = ent.Agent.Summary
						}
					case "acknowledge_log_entry":
						ts, err := time.Parse(time.RFC3339, ent.CreatedAt)
						if err != nil {
							return err
						}
						acknowledgeUnix = ts.Unix()
						acknowledgeUTC = ts.UTC()
					default:
						continue
					}
				}
			default:
				if len(incident.Acknowledgements) > 0 && len(incident.Acknowledgements[0].At) > 0 {
					ts, err := time.Parse(time.RFC3339, incident.Acknowledgements[0].At)
					if err != nil {
						return err
					}
					acknowledgeUnix = ts.Unix()
					acknowledgeUTC = ts.UTC()
				}
				if len(incident.Assignments) > 0 {
					assignTo = incident.Assignments[len(incident.Assignments)-1].Assignee.Summary
				}
			}

			service = incident.Service.Summary
			if len(incident.Teams) > 0 {
				team = incident.Teams[0].Summary
			}
			ePolicy = incident.EscalationPolicy.Summary
			if incident.Priority != nil {
				priority = incident.Priority.Name
			}
			for _, a := range alerts.Alerts {
				var triggerUnix int64 = 0
				var resolveUnix int64 = 0
				var triggerUTC time.Time
				var resolveUTC time.Time
				if ts, err := time.Parse(time.RFC3339, a.CreatedAt); err == nil {
					triggerUnix = ts.Unix()
					triggerUTC = ts.UTC()
				}
				if ts, err := time.Parse(time.RFC3339, a.ResolvedAt); err == nil {
					resolveUnix = ts.Unix()
					resolveUTC = ts.UTC()
				}
				alert := Alert{
					AlertID:          a.ID,
					IncidentID:       a.Incident.ID,
					TriggeredAt:      triggerUTC.Format(flinkTSFormat),
					TriggerAtUnix:    triggerUnix,
					Severity:         a.Severity,
					Urgency:          incident.Urgency,
					Priority:         priority,
					Service:          service,
					Team:             team,
					AssignTo:         assignTo,
					EscalationPolicy: ePolicy,
				}
				if acknowledgeUnix > 0 {
					alert.AcknowledgedAt = acknowledgeUTC.Format(flinkTSFormat)
					alert.AcknowledgedAtUnix = acknowledgeUnix
				}
				if resolveUnix > 0 {
					alert.ResolvedAt = resolveUTC.Format(flinkTSFormat)
					alert.ResolveAtUnix = resolveUnix
				}
				tid, cid, name, err := extract(a.Body)
				if err != nil {
					log.Printf("incident %v, alert %v, %v\n", a.Incident.ID, a.ID, err)
				}
				alert.TenantID, alert.ClusterID, alert.Name = tid, cid, name
				c.sink <- alert
			}
		}
	}
	return nil
}

func (c *collector) collect(ctx context.Context) error {
	defer close(c.sink)

	// collect yesterday resolved alert.
	{
		opts := new(pagerduty.ListIncidentsOptions)
		opts.Limit = 90
		opts.Offset = 0
		opts.Since = c.st.Format(time.RFC3339)
		opts.Until = c.et.Format(time.RFC3339)
		opts.SortBy = "created_at"
		opts.TimeZone = "UTC"
		opts.Statuses = []string{"resolved"}
		if err := c.collectAlert(ctx, *opts); err != nil {
			log.Println(err)
			return err
		}
	}
	// collect all triggered or acknowledged alert.
	{
		opts := new(pagerduty.ListIncidentsOptions)
		opts.Limit = 90
		opts.Offset = 0
		opts.DateRange = "all"
		opts.SortBy = "created_at"
		opts.TimeZone = "UTC"
		opts.Statuses = []string{"triggered", "acknowledged"}
		if err := c.collectAlert(ctx, *opts); err != nil {
			log.Println(err)
			return err
		}
	}
	return nil
}

// drain collect all alert
func (c *collector) drain(ctx context.Context) error {
	defer close(c.sink)
	opts := new(pagerduty.ListIncidentsOptions)
	opts.Limit = 500
	opts.Offset = 0
	opts.Since = c.drainSince
	opts.Until = c.drainUntil
	opts.SortBy = "created_at"
	opts.TimeZone = "UTC"
	if err := c.collectAlert(ctx, *opts); err != nil {
		log.Println(err)
		return err
	}
	return nil
}

func parseLabel(s string) (tenantID string, clusterID string, alertName string, err error) {
	s = strings.Split(s, "Annotations:")[0]
	// for corner case
	s = strings.ReplaceAll(s, ": ", ":")
	s = strings.ReplaceAll(s, "\n\t", " ")
	data := new(struct{
		Labels []string `yaml:"Labels"`
	})
	kvs := make(map[string]string)

	if err = yaml.Unmarshal([]byte(s), data); err != nil{
		return
	}
	for _, label := range data.Labels {
		if len(label) == 0 {
			continue
		}
		pair := strings.Split(string(label), "=")
		if len(pair) != 2 {
			continue
		}
		pair[0] = strings.TrimSpace(pair[0])
		pair[1] = strings.TrimSpace(pair[1])
		if pair[0] == "kubernetes_namespace" || pair[0] == "cluster_id" {
			pair[1] = strings.TrimLeft(strings.ToLower(pair[1]), "tidb")
			// for case cluster_id = tidb1311285575352848384/ab-1311285575352848384-2021-01-09t18-00-00
			pair[1] = strings.Split(pair[1], "/")[0]
		}
		kvs[pair[0]] = pair[1]
	}
	alertName = kvs["alertname"]
	if v, ok := kvs["tenant"]; ok {
		tenantID = v
	}else {
		tenantID = kvs["tenant_id"]
	}
	if v, ok := kvs["cluster_id"]; ok {
		clusterID = v
	}else {
		clusterID = kvs["kubernetes_namespace"]
	}
	return
}

func extract(body map[string]interface{}) (string, string, string, error) {
	details, ok := body["details"].(map[string]interface{})
	if !ok {
		return "", "", "", nil
	}
	s, ok := details["firing"].(string)
	if !ok {
		return "", "", "", nil
	}
	return parseLabel(s)
}