package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"strconv"
	"time"

	"github.com/cenkalti/backoff/v4"
	"github.com/michimani/gotwi"
	"github.com/michimani/gotwi/tweet/managetweet"
	"github.com/michimani/gotwi/tweet/managetweet/types"
	"go.skia.org/infra/go/monorail/v1"
	"golang.org/x/exp/slices"
	"google.golang.org/api/option"
)

const (
	monorailEndpointURL = "https://bugs.avm99963.com/_ah/api/monorail/v1/"
	vulnzProject        = "vulnz"
	doNotPublishLabel   = "DoNotPublish"
	restrictedLabel     = "Restrict-View-Commit"
	// Days after which the report for a fixed vulnerability will be automatically
	// published.
	deadlineAfterFixed = 30 * 24 * time.Hour
	// Grace period added to the deadlines
	gracePeriod = 1 * 24 * time.Hour
	// Time to warn about the automatic disclosure before it happens
	warningPeriod = 5 * 24 * time.Hour
	// Layouts returned by the monorail package
	dateTimeLayout = "2006-01-02T15:04:05"
	dateLayout     = "2006-01-02"
)

type VulnerabilityReport struct {
	// ID of the vulnerability report/issue in Monorail
	IssueId int64
	// Title of the report
	Title string
	// Issue status
	Status string
	// Date that the status was modified
	StatusModified time.Time
	// Whether a PublishAt field exists
	HasPublishAt bool
	// Custom date when the vulnerability should be published. Takes precedence over the default disclosure flow
	PublishAt time.Time
	// Date when the vulnerability report was sent to the vendor
	Reported time.Time
	// Number of days since the reported date when the vulnerability should be automatically disclosed
	DeadlineDays int
	// DeadlineDays, converted to time.Duration
	DeadlineDuration time.Duration
	// Date when the vulnerability report should be published if it has been fixed
	PublishDateIfFixed time.Time
	// Whether DeadlineTime could be calculated or not. If false, the vulnerability should not be published if not fixed
	HasValidDeadlineTime bool
	// Deadline when the vulnerability report should be published if not fixed
	DeadlineTime         time.Time
	HasDoNotPublishLabel bool
	HasRestrictedLabel   bool
}

type Action struct {
	MustPerform bool
	// |Reason| can be `disclosureDeadline` (the deadline has been exceeded), `fixedDeadline` (the issue has been fixed and |deadlineAfterFixed| has passed), `customDisclosureDeadline` (the custom PublishAt date has passed)
	Reason string
	// |Type| can be `warning` (warning before actually disclosing the report), `disclosure` (the report is actually disclosed)
	Type string
}

func parseIssue(issue *monorail.ProtoApiPb2V1IssueWrapper) (*VulnerabilityReport, error) {
	report := VulnerabilityReport{IssueId: issue.Id}

	report.Title = issue.Title
	report.Status = issue.Status

	var err error
	report.StatusModified, err = time.Parse(dateTimeLayout, issue.StatusModified)
	if err != nil {
		return nil, err
	}

	deadlineFound := false
	reportedFound := false
	report.HasPublishAt = false
	for _, field := range issue.FieldValues {
		if field.FieldName == "Deadline" {
			deadlineFound = true
			report.DeadlineDays, err = strconv.Atoi(field.FieldValue)
			report.DeadlineDuration = time.Duration(report.DeadlineDays) * 24 * time.Hour
			if err != nil {
				return nil, fmt.Errorf("Error parsing Deadline field: %v")
			}
		}

		if field.FieldName == "Reported" {
			reportedFound = true
			report.Reported, err = time.Parse(dateLayout, field.FieldValue)
			if err != nil {
				return nil, fmt.Errorf("Error parsing Reported field: %v")
			}
		}

		if field.FieldName == "PublishAt" {
			report.PublishAt, err = time.Parse(dateLayout, field.FieldValue)
			if err != nil {
				return nil, fmt.Errorf("Error parsing PublishAt field: %v")
			}
			report.PublishAt = report.PublishAt.Add(gracePeriod)
			report.HasPublishAt = true
		}
	}

	report.PublishDateIfFixed = report.StatusModified.Add(deadlineAfterFixed + gracePeriod)
	if reportedFound && deadlineFound {
		report.HasValidDeadlineTime = true
		report.DeadlineTime = report.Reported.Add(report.DeadlineDuration + gracePeriod)
	} else {
		report.HasValidDeadlineTime = false
	}

	report.HasDoNotPublishLabel = slices.Contains(issue.Labels, doNotPublishLabel)
	report.HasRestrictedLabel = slices.Contains(issue.Labels, restrictedLabel)

	return &report, nil
}

func logIssue(message string, report *VulnerabilityReport) {
	log.Printf("[Issue %d] %s", report.IssueId, message)
}

// |actionType| can be one of the types allowed in |Action.Type|.
func getActionPartial(report *VulnerabilityReport, actionType string) *Action {
	now := time.Now()
	isWarning := actionType == "warning"

	// Take into account that warnings are published in advance (according to warningPeriod)
	var (
		customPublishAtDate time.Time
		publishDateIfFixed  time.Time
		deadlineTime        time.Time
	)
	if isWarning {
		customPublishAtDate = report.PublishAt.Add(-warningPeriod)
		publishDateIfFixed = report.PublishDateIfFixed.Add(-warningPeriod)
		deadlineTime = report.DeadlineTime.Add(-warningPeriod)
	} else {
		customPublishAtDate = report.PublishAt
		publishDateIfFixed = report.PublishDateIfFixed
		deadlineTime = report.DeadlineTime
	}

	if report.HasPublishAt {
		if (!isWarning && customPublishAtDate.Before(now)) || (isWarning && dateEqual(customPublishAtDate, now)) {
			return &Action{
				MustPerform: true,
				Reason:      "customDisclosureDeadline",
				Type:        actionType,
			}
		}
		return &Action{MustPerform: false}
	}
	if report.Status == "Fixed" || report.Status == "Verified" {
		if (!isWarning && publishDateIfFixed.Before(now)) || (isWarning && dateEqual(publishDateIfFixed, now)) {
			return &Action{
				MustPerform: true,
				Reason:      "fixedDeadline",
				Type:        actionType,
			}
		}
		return &Action{MustPerform: false}
	}
	if report.HasValidDeadlineTime && (!isWarning && deadlineTime.Before(now)) || (isWarning && dateEqual(deadlineTime, now)) {
		return &Action{
			MustPerform: true,
			Reason:      "disclosureDeadline",
			Type:        actionType,
		}
	}
	return &Action{MustPerform: false}
}

// Get the actions that should be performed today on |report|
func getTodayActions(report *VulnerabilityReport) []*Action {
	// By default, all the actions will be set to |false|
	var actions []*Action

	// Do not publish any issue which has the DoNotPublish label set!
	if report.HasDoNotPublishLabel {
		logIssue(fmt.Sprintf("The %s label is set. Skipping.\n", doNotPublishLabel), report)
		return actions
	}

	// Check that the report is not already published
	if !report.HasRestrictedLabel {
		logIssue(fmt.Sprintf("The %s label isn't set (so in theory it shouldn't have been returned by the query). Skipping.\n", restrictedLabel), report)
		return actions
	}

	warningAction := getActionPartial(report, "warning")
	if warningAction.MustPerform {
		actions = append(actions, warningAction)
	}
	disclosureAction := getActionPartial(report, "disclosure")
	if disclosureAction.MustPerform {
		actions = append(actions, disclosureAction)
	}
	return actions
}

func newRequestExponentialBackOff() *backoff.ExponentialBackOff {
	b := backoff.NewExponentialBackOff()
	b.InitialInterval = 500 * time.Millisecond
	b.RandomizationFactor = 0.5
	b.Multiplier = 1.5
	b.MaxInterval = 30 * time.Second
	b.MaxElapsedTime = 5 * time.Minute
	b.Reset()
	return b
}

func publishReportAttempt(message string, report *VulnerabilityReport, action *Action, monorailSvc *monorail.Service, ctx context.Context) error {
	ctxRequest, cancel := context.WithDeadline(ctx, time.Now().Add(45*time.Second))
	defer cancel()

	comment := monorail.ProtoApiPb2V1IssueCommentWrapper{
		Content: message,
	}

	if action.Type == "disclosure" {
		// Remove the label which restricts the bug from public view
		labels := make([]string, 1)
		labels[0] = "-" + restrictedLabel
		comment.Updates = &monorail.ProtoApiPb2V1Update{Labels: labels}
	}

	_, err := monorailSvc.Issues.Comments.Insert(vulnzProject, report.IssueId, &comment).SendEmail(true).Context(ctxRequest).Do()
	if err != nil {
		return fmt.Errorf("The request to insert the comment failed: %v", err)
	}
	logIssue(fmt.Sprintf("Added comment (type: %s, reason: %s)", action.Type, action.Reason), report)

	return nil
}

func publishReport(message string, report *VulnerabilityReport, action *Action, monorailSvc *monorail.Service, ctx context.Context) error {
	exp := newRequestExponentialBackOff()
	err := backoff.Retry(func() error {
		return publishReportAttempt(message, report, action, monorailSvc, ctx)
	}, exp)
	return err
}

func postTweet(report *VulnerabilityReport, twtrClient *gotwi.Client, ctx context.Context) error {
	truncatedBugTitle := truncateStringWithEllipsis(report.Title, 240-(2+1+6+1+1+20))
	url := fmt.Sprintf("iavm.xyz/b/vulnz/%d", report.IssueId)
	tweetText := fmt.Sprintf("🦋 b/%d %s %s", report.IssueId, truncatedBugTitle, url)
	body := &types.CreateInput{
		Text: gotwi.String(tweetText),
	}
	_, err := managetweet.Create(ctx, twtrClient, body)
	if err == nil {
		logIssue(fmt.Sprintf("Tweeted: %s", tweetText), report)
	}
	return err
}

func performAction(report *VulnerabilityReport, action *Action, monorailSvc *monorail.Service, twtrClient *gotwi.Client, ctx context.Context) error {
	var message string
	if action.Reason == "fixedDeadline" {
		deadlineAfterFixedInt := deadlineAfterFixed / (24 * time.Hour)
		if action.Type == "warning" {
			message = fmt.Sprintf("This report will be published in %d days, since at that time %d days will have passed since the vulnerability was fixed.\n\n_Please add the `%s` label if you want to stop the automatic disclosure._", warningPeriod/(24*time.Hour), deadlineAfterFixedInt, doNotPublishLabel)
		} else {
			message = fmt.Sprintf("**The vulnerability was fixed %d days ago** -- automatically publishing the vulnerability report.", deadlineAfterFixedInt)
		}
	} else if action.Reason == "disclosureDeadline" {
		if action.Type == "warning" {
			message = fmt.Sprintf("This report will be published in %d days, since at that time %d days will have passed since the vulnerability was reported.\n\n_Please add the `%s` label if you want to stop the automatic disclosure._", warningPeriod/(24*time.Hour), report.DeadlineDays, doNotPublishLabel)
		} else {
			message = fmt.Sprintf("**The %d-day deadline has been exceeded** -- automatically publishing the vulnerability report.", report.DeadlineDays)
		}
	} else if action.Reason == "customDisclosureDeadline" {
		if action.Type == "warning" {
			message = fmt.Sprintf("This report will be published in %d days, as set in the PublishAt field.\n\n_Please unset `PublishAt` or add the `%s` label if you want to stop the automatic disclosure._", warningPeriod/(24*time.Hour), doNotPublishLabel)
		} else {
			message = "**The PublishAt date has passed** -- automatically publishing the vulnerability report."
		}
	}
	err := publishReport(message, report, action, monorailSvc, ctx)
	if err != nil {
		return fmt.Errorf("Error trying to perform action type %s, for reason %s: %v", action.Type, action.Reason, err)
	}

	if twtrClient != nil && action.Type == "disclosure" {
		err := postTweet(report, twtrClient, ctx)
		if err != nil {
			return fmt.Errorf("Error trying to publish tweet: %v", err)
		}
	}

	return nil
}

func handleIssue(issue *monorail.ProtoApiPb2V1IssueWrapper, monorailSvc *monorail.Service, twtrClient *gotwi.Client, ctx context.Context) error {
	report, err := parseIssue(issue)
	if err != nil {
		return fmt.Errorf("Error parsing issue: %v", err)
	}

	actions := getTodayActions(report)

	for _, action := range actions {
		err = performAction(report, action, monorailSvc, twtrClient, ctx)
		if err != nil {
			return fmt.Errorf("Error handling %s action: %v", action.Type, err)
		}
	}

	return nil
}

func main() {
	log.Println("Starting vulnzybot...")

	ctx := context.Background()

	credentials, ok := os.LookupEnv("GOOGLE_APPLICATION_CREDENTIALS")
	if !ok {
		log.Println("GOOGLE_APPLICATION_CREDENTIALS environment variable is not set.")
		return
	}

	monorailSvc, err := monorail.NewService(ctx, option.WithEndpoint(monorailEndpointURL), option.WithCredentialsFile(credentials))
	if err != nil {
		log.Fatalf("Can't create Monorail service: %v", err)
	}

	twtrToken, okToken := os.LookupEnv("TWITTER_OAUTH_TOKEN")
	twtrSecret, okSecret := os.LookupEnv("TWITTER_OAUTH_TOKEN_SECRET")
	var twtrClient *gotwi.Client
	if okToken && okSecret {
		in := &gotwi.NewClientInput{
			AuthenticationMethod: gotwi.AuthenMethodOAuth1UserContext,
			OAuthToken:           twtrToken,
			OAuthTokenSecret:     twtrSecret,
		}
		twtrClient, err = gotwi.NewClient(in)
		if err != nil {
			log.Fatalf("Can't create Twitter client: %v", err)
		}
	}

	var listResponse *monorail.ProtoApiPb2V1IssuesListResponse

	exp := newRequestExponentialBackOff()
	err = backoff.Retry(func() error {
		var err error
		listResponse, err = monorailSvc.Issues.List(vulnzProject).Can("all").Q("Type=VulnerabilityReport Restrict=View-Commit").MaxResults(100000).Do()
		return err
	}, exp)
	if err != nil {
		log.Fatalf("Can't list issues: %v", err)
	}

	for _, issue := range listResponse.Items {
		err := handleIssue(issue, monorailSvc, twtrClient, ctx)
		if err != nil {
			log.Printf("[Issue %v] Error while handling issue: %v", issue.Id, err)
		}
	}
}