package GCP_CLOUD_FUNCTIONS

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	securitycenter "cloud.google.com/go/securitycenter/apiv1"
	"cloud.google.com/go/securitycenter/apiv1/securitycenterpb"
)

const (
	GCP_Folder_ID      = ""
	AuthorizationToken = ""
)

// parseTime parses the time parameter from header, supporting both hours (Nh) and days (Nd)
// Examples: "6h" for 6 hours, "1d" for 1 day, "12" for 12 days (backward compatibility)
func parseTime(timeStr string) time.Time {
	if timeStr == "" {
		// default to 1 day
		return time.Now().UTC().AddDate(0, 0, -1)
	}

	// Check if it ends with 'h' for hours
	if strings.HasSuffix(timeStr, "h") {
		hoursStr := strings.TrimSuffix(timeStr, "h")
		hours, err := strconv.Atoi(hoursStr)
		if err != nil || hours <= 0 {
			// default to 1 day if invalid
			return time.Now().UTC().AddDate(0, 0, -1)
		}
		return time.Now().UTC().Add(-time.Duration(hours) * time.Hour)
	}

	// Check if it ends with 'd' for days
	if strings.HasSuffix(timeStr, "d") {
		daysStr := strings.TrimSuffix(timeStr, "d")
		days, err := strconv.Atoi(daysStr)
		if err != nil || days <= 0 {
			// default to 1 day if invalid
			return time.Now().UTC().AddDate(0, 0, -1)
		}
		return time.Now().UTC().AddDate(0, 0, -days)
	}

	// Backward compatibility: treat as days if no suffix
	days, err := strconv.Atoi(timeStr)
	if err != nil || days <= 0 {
		// default to 1 day if invalid
		return time.Now().UTC().AddDate(0, 0, -1)
	}
	return time.Now().UTC().AddDate(0, 0, -days)
}

// parseSeverity parses the severity levels from header or uses default
func parseSeverity(severityStr string) []string {
	if severityStr == "" {
		return []string{"CRITICAL"} // default to CRITICAL
	}

	// Split by comma and trim spaces
	severityList := strings.Split(severityStr, ",")
	var result []string
	for _, s := range severityList {
		trimmed := strings.TrimSpace(strings.ToUpper(s))
		if trimmed != "" {
			result = append(result, trimmed)
		}
	}

	// If no valid severity levels, default to CRITICAL
	if len(result) == 0 {
		return []string{"CRITICAL"}
	}

	return result
}

// buildSeverityFilter builds the severity filter string
func buildSeverityFilter(severities []string) string {
	if len(severities) == 0 {
		return `severity="CRITICAL"`
	}

	if len(severities) == 1 {
		return fmt.Sprintf(`severity="%s"`, severities[0])
	}

	var parts []string
	for _, s := range severities {
		parts = append(parts, fmt.Sprintf(`severity="%s"`, s))
	}
	return strings.Join(parts, " OR ")
}

// parsePageSize parses the page size from header or uses default
func parsePageSize(pageSizeStr string) int32 {
	if pageSizeStr == "" {
		return 100 // default to 100
	}
	pageSize, err := strconv.Atoi(pageSizeStr)
	if err != nil || pageSize <= 0 {
		return 100 // default to 100 if invalid
	}
	// Maximum page size is 1000 according to API documentation
	if pageSize > 1000 {
		return 1000
	}
	return int32(pageSize)
}

// ApiQuery_THREAT handles the THREAT findings query
func ApiQuery_THREAT(w http.ResponseWriter, r *http.Request) {
	// Set response header to JSON
	w.Header().Set("Content-Type", "application/json")

	// Create a context
	ctx := context.Background()

	// Create a client for Security Command Center
	client, err := securitycenter.NewClient(ctx)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create SCC client: %v", err), http.StatusInternalServerError)
		return
	}
	defer client.Close()

	// Get page size and page token from headers
	pageSize := parsePageSize(r.Header.Get("Page_size"))
	pageToken := r.Header.Get("Page_token")

	// Get filter from header if provided (for pagination requests)
	filterFromHeader := r.Header.Get("Filter")
	var filter string

	if filterFromHeader != "" {
		// Decode the base64 encoded filter from header for pagination requests
		decodedFilter, err := base64.StdEncoding.DecodeString(filterFromHeader)
		if err != nil {
			http.Error(w, fmt.Sprintf("Failed to decode filter: %v", err), http.StatusBadRequest)
			return
		}
		filter = string(decodedFilter)
	} else {
		// Calculate filter parameters for initial requests
		startTime := parseTime(r.Header.Get("Time")).Format("2006-01-02T15:04:05Z")

		// Get severity levels from header or use default
		severities := parseSeverity(r.Header.Get("Servity")) // Note: Keeping the typo from user request
		severityFilter := buildSeverityFilter(severities)

		// Build the filter string
		filter = fmt.Sprintf(`state="ACTIVE" AND NOT mute="MUTED" AND finding_class="THREAT" AND (%s) AND event_time>="%s"`, severityFilter, startTime)
	}

	// Replace with your folder ID
	folderID := GCP_Folder_ID

	// Create the request to list findings with date and severity filtering
	// The parent is the folder resource
	req := &securitycenterpb.ListFindingsRequest{
		Parent:    fmt.Sprintf("folders/%s/sources/-", folderID),
		Filter:    filter,
		PageSize:  pageSize,
		PageToken: pageToken,
	}

	// Execute the request
	it := client.ListFindings(ctx, req)

	// Collect findings with specific fields like in the gcloud command
	type FindingSummary struct {
		Category         string `json:"category"`
		Severity         string `json:"severity"`
		AffectedResource string `json:"affected_resource"`
		ResourceApiName  string `json:"resource_api_name"`
		EventTime        string `json:"event_time"`
		CreateTime       string `json:"create_time"`
	}

	var findings []FindingSummary
	var nextPageToken string

	// Use the iterator's internal fetch mechanism to get one page of results
	// This will properly handle the page token and return the next page token
	items, token, err := it.InternalFetch(int(pageSize), pageToken)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to fetch findings: %v", err), http.StatusInternalServerError)
		return
	}

	// Process the items we got
	for _, result := range items {
		// Extract the finding from the result
		finding := result.Finding
		resource := result.Resource

		// Add to our findings slice
		findings = append(findings, FindingSummary{
			Category:         finding.GetCategory(),
			Severity:         finding.GetSeverity().String(),
			AffectedResource: resource.GetDisplayName(),
			ResourceApiName:  resource.GetName(),
			EventTime:        finding.GetEventTime().AsTime().Format(time.RFC3339),
			CreateTime:       finding.GetCreateTime().AsTime().Format(time.RFC3339),
		})
	}

	// Set the next page token
	nextPageToken = token

	// Base64 encode the filter for safe transmission
	encodedFilter := base64.StdEncoding.EncodeToString([]byte(filter))

	// Convert findings to JSON and send as response
	jsonData, err := json.Marshal(findings)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to marshal findings to JSON: %v", err), http.StatusInternalServerError)
		return
	}

	// Send the response with next page token and base64 encoded filter
	fmt.Fprintf(w, `{"findings_count": %d, "next_page_token": "%s", "filter_used": "%s", "findings": %s}`, len(findings), nextPageToken, encodedFilter, string(jsonData))
}

// ApiQuery_TOXIC_COMBINATION handles the TOXIC_COMBINATION findings query
func ApiQuery_TOXIC_COMBINATION(w http.ResponseWriter, r *http.Request) {
	// Set response header to JSON
	w.Header().Set("Content-Type", "application/json")

	// Create a context
	ctx := context.Background()

	// Create a client for Security Command Center
	client, err := securitycenter.NewClient(ctx)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create SCC client: %v", err), http.StatusInternalServerError)
		return
	}
	defer client.Close()

	// Get page size and page token from headers
	pageSize := parsePageSize(r.Header.Get("Page_size"))
	pageToken := r.Header.Get("Page_token")

	// Get filter from header if provided (for pagination requests)
	filterFromHeader := r.Header.Get("Filter")
	var filter string

	if filterFromHeader != "" {
		// Decode the base64 encoded filter from header for pagination requests
		decodedFilter, err := base64.StdEncoding.DecodeString(filterFromHeader)
		if err != nil {
			http.Error(w, fmt.Sprintf("Failed to decode filter: %v", err), http.StatusBadRequest)
			return
		}
		filter = string(decodedFilter)
	} else {
		// Calculate filter parameters for initial requests
		startTime := parseTime(r.Header.Get("Time")).Format("2006-01-02T15:04:05Z")

		// Get severity levels from header or use default
		severities := parseSeverity(r.Header.Get("Servity")) // Note: Keeping the typo from user request
		severityFilter := buildSeverityFilter(severities)

		// Build the filter string
		filter = fmt.Sprintf(`state="ACTIVE" AND NOT mute="MUTED" AND finding_class="TOXIC_COMBINATION" AND (%s) AND event_time>="%s"`, severityFilter, startTime)
	}

	// Replace with your folder ID
	folderID := GCP_Folder_ID

	// Create the request to list findings with date and severity filtering
	// The parent is the folder resource
	req := &securitycenterpb.ListFindingsRequest{
		Parent:    fmt.Sprintf("folders/%s/sources/-", folderID),
		Filter:    filter,
		PageSize:  pageSize,
		PageToken: pageToken,
	}

	// Execute the request
	it := client.ListFindings(ctx, req)

	// Collect findings with specific fields like in the gcloud command
	type FindingSummary struct {
		Category            string  `json:"category"`
		Severity            string  `json:"severity"`
		AttackExposureScore float64 `json:"attack_exposure_score"`
		AffectedResource    string  `json:"affected_resource"`
		ResourceApiName     string  `json:"resource_api_name"`
		EventTime           string  `json:"event_time"`
		CreateTime          string  `json:"create_time"`
	}

	var findings []FindingSummary
	var nextPageToken string

	// Use the iterator's internal fetch mechanism to get one page of results
	// This will properly handle the page token and return the next page token
	items, token, err := it.InternalFetch(int(pageSize), pageToken)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to fetch findings: %v", err), http.StatusInternalServerError)
		return
	}

	// Process the items we got
	for _, result := range items {
		// Extract the finding from the result
		finding := result.Finding
		resource := result.Resource

		// Extract attack exposure score if available
		var attackExposureScore float64
		if finding.GetToxicCombination() != nil {
			attackExposureScore = finding.GetToxicCombination().GetAttackExposureScore()
		}

		// Add to our findings slice
		findings = append(findings, FindingSummary{
			Category:            finding.GetCategory(),
			Severity:            finding.GetSeverity().String(),
			AttackExposureScore: attackExposureScore,
			AffectedResource:    resource.GetDisplayName(),
			ResourceApiName:     resource.GetName(),
			EventTime:           finding.GetEventTime().AsTime().Format(time.RFC3339),
			CreateTime:          finding.GetCreateTime().AsTime().Format(time.RFC3339),
		})
	}

	// Set the next page token
	nextPageToken = token

	// Base64 encode the filter for safe transmission
	encodedFilter := base64.StdEncoding.EncodeToString([]byte(filter))

	// Convert findings to JSON and send as response
	jsonData, err := json.Marshal(findings)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to marshal findings to JSON: %v", err), http.StatusInternalServerError)
		return
	}

	// Send the response with next page token and base64 encoded filter
	fmt.Fprintf(w, `{"findings_count": %d, "next_page_token": "%s", "filter_used": "%s", "findings": %s}`, len(findings), nextPageToken, encodedFilter, string(jsonData))
}

// ApiQuery_RISK_BULLETIN handles the RISK BULLETIN findings query
func ApiQuery_RISK_BULLETIN(w http.ResponseWriter, r *http.Request) {
	// Set response header to JSON
	w.Header().Set("Content-Type", "application/json")

	// Create a context
	ctx := context.Background()

	// Create a client for Security Command Center
	client, err := securitycenter.NewClient(ctx)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create SCC client: %v", err), http.StatusInternalServerError)
		return
	}
	defer client.Close()

	// Get page size and page token from headers
	pageSize := parsePageSize(r.Header.Get("Page_size"))
	pageToken := r.Header.Get("Page_token")

	// Get filter from header if provided (for pagination requests)
	filterFromHeader := r.Header.Get("Filter")
	var filter string

	if filterFromHeader != "" {
		// Decode the base64 encoded filter from header for pagination requests
		decodedFilter, err := base64.StdEncoding.DecodeString(filterFromHeader)
		if err != nil {
			http.Error(w, fmt.Sprintf("Failed to decode filter: %v", err), http.StatusBadRequest)
			return
		}
		filter = string(decodedFilter)
	} else {
		// Calculate filter parameters for initial requests
		startTime := parseTime(r.Header.Get("Time")).Format("2006-01-02T15:04:05Z")

		// Get severity levels from header or use default
		severities := parseSeverity(r.Header.Get("Servity")) // Note: Keeping the typo from user request
		severityFilter := buildSeverityFilter(severities)

		// Build the filter string for risk bulletin findings
		filter = fmt.Sprintf(`state="ACTIVE" AND NOT mute="MUTED" AND category="GKE_SECURITY_BULLETIN" AND (%s) AND event_time>="%s"`, severityFilter, startTime)
	}

	// Replace with your folder ID
	folderID := GCP_Folder_ID

	// Create the request to list findings with date and severity filtering
	// The parent is the folder resource
	req := &securitycenterpb.ListFindingsRequest{
		Parent:    fmt.Sprintf("folders/%s/sources/-", folderID),
		Filter:    filter,
		PageSize:  pageSize,
		PageToken: pageToken,
	}

	// Execute the request
	it := client.ListFindings(ctx, req)

	// Collect findings with specific fields like in the gcloud command
	type FindingSummary struct {
		Category         string `json:"category"`
		BulletinId       string `json:"bulletin_id"`
		UpgradeVersion   string `json:"upgrade_version"`
		Severity         string `json:"severity"`
		Project          string `json:"project"`
		Location         string `json:"location"`
		NodePool         string `json:"nodepool"`
		AffectedResource string `json:"affected_resource"`
		Description      string `json:"description"`
		EventTime        string `json:"event_time"`
		CreateTime       string `json:"create_time"`
	}

	var findings []FindingSummary
	var nextPageToken string

	// Use the iterator's internal fetch mechanism to get one page of results
	// This will properly handle the page token and return the next page token
	items, token, err := it.InternalFetch(int(pageSize), pageToken)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to fetch findings: %v", err), http.StatusInternalServerError)
		return
	}

	// Process the items we got
	for _, result := range items {
		// Extract the finding from the result
		finding := result.Finding
		resource := result.Resource

		// Extract bulletin information if available
		var bulletinId, upgradeVersion string
		if finding.GetVulnerability() != nil && finding.GetVulnerability().GetSecurityBulletin() != nil {
			bulletinId = finding.GetVulnerability().GetSecurityBulletin().GetBulletinId()
			upgradeVersion = finding.GetVulnerability().GetSecurityBulletin().GetSuggestedUpgradeVersion()
		}

		// Extract node pool information if available
		var nodePool string
		if len(finding.GetKubernetes().GetNodePools()) > 0 {
			nodePool = finding.GetKubernetes().GetNodePools()[0].GetName()
		}

		// Add to our findings slice
		findings = append(findings, FindingSummary{
			Category:         finding.GetCategory(),
			BulletinId:       bulletinId,
			UpgradeVersion:   upgradeVersion,
			Severity:         finding.GetSeverity().String(),
			Project:          resource.GetProjectDisplayName(),
			Location:         resource.GetLocation(),
			NodePool:         nodePool,
			AffectedResource: resource.GetName(),
			Description:      finding.GetDescription(),
			EventTime:        finding.GetEventTime().AsTime().Format(time.RFC3339),
			CreateTime:       finding.GetCreateTime().AsTime().Format(time.RFC3339),
		})
	}

	// Set the next page token
	nextPageToken = token

	// Base64 encode the filter for safe transmission
	encodedFilter := base64.StdEncoding.EncodeToString([]byte(filter))

	// Convert findings to JSON and send as response
	jsonData, err := json.Marshal(findings)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to marshal findings to JSON: %v", err), http.StatusInternalServerError)
		return
	}

	// Send the response with next page token and base64 encoded filter
	fmt.Fprintf(w, `{"findings_count": %d, "next_page_token": "%s", "filter_used": "%s", "findings": %s}`, len(findings), nextPageToken, encodedFilter, string(jsonData))
}

// ApiQuery_OS_VULNERABILITY handles the OS VULNERABILITY findings query
func ApiQuery_OS_VULNERABILITY(w http.ResponseWriter, r *http.Request) {
	// Set response header to JSON
	w.Header().Set("Content-Type", "application/json")

	// Create a context
	ctx := context.Background()

	// Create a client for Security Command Center
	client, err := securitycenter.NewClient(ctx)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create SCC client: %v", err), http.StatusInternalServerError)
		return
	}
	defer client.Close()

	// Get page size and page token from headers
	pageSize := parsePageSize(r.Header.Get("Page_size"))
	pageToken := r.Header.Get("Page_token")

	// Get filter from header if provided (for pagination requests)
	filterFromHeader := r.Header.Get("Filter")
	var filter string

	if filterFromHeader != "" {
		// Decode the base64 encoded filter from header for pagination requests
		decodedFilter, err := base64.StdEncoding.DecodeString(filterFromHeader)
		if err != nil {
			http.Error(w, fmt.Sprintf("Failed to decode filter: %v", err), http.StatusBadRequest)
			return
		}
		filter = string(decodedFilter)
	} else {
		// Calculate filter parameters for initial requests
		startTime := parseTime(r.Header.Get("Time")).Format("2006-01-02T15:04:05Z")

		// Get severity levels from header or use default
		severities := parseSeverity(r.Header.Get("Servity")) // Note: Keeping the typo from user request
		severityFilter := buildSeverityFilter(severities)

		// Build the filter string for OS vulnerability findings
		filter = fmt.Sprintf(`state="ACTIVE" AND NOT mute="MUTED" AND category="OS_VULNERABILITY" AND (%s) AND event_time>="%s"`, severityFilter, startTime)
	}

	// Replace with your folder ID
	folderID := GCP_Folder_ID

	// Create the request to list findings with date and severity filtering
	// The parent is the folder resource
	req := &securitycenterpb.ListFindingsRequest{
		Parent:    fmt.Sprintf("folders/%s/sources/-", folderID),
		Filter:    filter,
		PageSize:  pageSize,
		PageToken: pageToken,
	}

	// Execute the request
	it := client.ListFindings(ctx, req)

	// Collect findings with specific fields like in the gcloud command
	type FindingSummary struct {
		Category            string `json:"category"`
		Severity            string `json:"severity"`
		AttackExposureState string `json:"attack_exposure_state"`
		EventTime           string `json:"event_time"`
		CreateTime          string `json:"create_time"`
		FindingClass        string `json:"finding_class"`
		AffectedResource    string `json:"affected_resource"`
		ResourcePath        string `json:"resource_path"`
	}

	var findings []FindingSummary
	var nextPageToken string

	// Use the iterator's internal fetch mechanism to get one page of results
	// This will properly handle the page token and return the next page token
	items, token, err := it.InternalFetch(int(pageSize), pageToken)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to fetch findings: %v", err), http.StatusInternalServerError)
		return
	}

	// Process the items we got
	for _, result := range items {
		// Extract the finding from the result
		finding := result.Finding
		resource := result.Resource

		// Extract attack exposure state if available
		var attackExposureState string
		if finding.GetAttackExposure() != nil {
			attackExposureState = finding.GetAttackExposure().GetState().String()
		}

		// Add to our findings slice
		findings = append(findings, FindingSummary{
			Category:            finding.GetCategory(),
			Severity:            finding.GetSeverity().String(),
			AttackExposureState: attackExposureState,
			EventTime:           finding.GetEventTime().AsTime().Format(time.RFC3339),
			CreateTime:          finding.GetCreateTime().AsTime().Format(time.RFC3339),
			FindingClass:        finding.GetFindingClass().String(),
			AffectedResource:    resource.GetDisplayName(),
			ResourcePath:        resource.GetResourcePathString(),
		})
	}

	// Set the next page token
	nextPageToken = token

	// Base64 encode the filter for safe transmission
	encodedFilter := base64.StdEncoding.EncodeToString([]byte(filter))

	// Convert findings to JSON and send as response
	jsonData, err := json.Marshal(findings)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to marshal findings to JSON: %v", err), http.StatusInternalServerError)
		return
	}

	// Send the response with next page token and base64 encoded filter
	fmt.Fprintf(w, `{"findings_count": %d, "next_page_token": "%s", "filter_used": "%s", "findings": %s}`, len(findings), nextPageToken, encodedFilter, string(jsonData))
}