package main

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

	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/service/sts"

	aws_collector "security-center-backend/internal/modules/caasm/collector/aws"
	caasmmodel "security-center-backend/internal/modules/caasm/model"
)

type CollectorEvent struct {
	SyncJobID      uint64   `json:"sync_job_id"`
	CloudAccountID uint     `json:"cloud_account_id"`
	RoleARN        string   `json:"role_arn"`
	ExternalID     string   `json:"external_id"`
	Region         string   `json:"region"`
	AllowedRegions []string `json:"allowed_regions"`
}

func HandleRequest(ctx context.Context, event CollectorEvent) (string, error) {
	log.Printf("Starting collection for account %d in region %s (Job %d)", event.CloudAccountID, event.Region, event.SyncJobID)

	// Load AWS Config
	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(event.Region))
	if err != nil {
		return "", fmt.Errorf("failed to load aws config: %v", err)
	}

	// Assume the target role using the Lambda execution role as source credentials
	if event.RoleARN != "" {
		stsClient := sts.NewFromConfig(cfg)
		provider := stscreds.NewAssumeRoleProvider(stsClient, event.RoleARN, func(o *stscreds.AssumeRoleOptions) {
			if event.ExternalID != "" {
				o.ExternalID = &event.ExternalID
			}
		})
		cfg.Credentials = aws.NewCredentialsCache(provider)
	}

	// Initialize Collector
	col, err := aws_collector.NewAWSCollectorWithConfig(cfg, event.Region, event.CloudAccountID, event.AllowedRegions)
	if err != nil {
		return "", fmt.Errorf("failed to create collector: %v", err)
	}

	// Collect
	result, err := col.Collect(ctx)
	if err != nil {
		return "", fmt.Errorf("collection failed: %v", err)
	}

	// Prepare Ingest Payload
	payload := caasmmodel.IngestPayload{
		SyncJobID:      event.SyncJobID,
		CloudAccountID: uint64(event.CloudAccountID),
		Assets:         result.Assets,
		IPAssets:       result.IPAssets,
		DomainAssets:   result.DomainAssets,
		Ports:          result.Ports,
	}

	payloadBytes, err := json.Marshal(payload)
	if err != nil {
		return "", fmt.Errorf("failed to marshal payload: %v", err)
	}

	// Send to Backend API
	apiURL := os.Getenv("API_URL")
	if apiURL == "" {
		// Fallback for local testing if needed, or error
		log.Println("API_URL not set, skipping ingestion (dry run)")
		return "Dry run complete", nil
	}

	// 1. Mark Job as Started
	if err := callLifecycleHook(ctx, apiURL, event.SyncJobID, "start", nil); err != nil {
		log.Printf("Failed to mark job as started: %v", err)
		// Continue anyway? Or fail? Let's continue but log.
	}

	// Defer failure reporting
	defer func() {
		if err != nil {
			failPayload := map[string]interface{}{
				"error_message": err.Error(),
			}
			_ = callLifecycleHook(ctx, apiURL, event.SyncJobID, "fail", failPayload)
		}
	}()

	// Report progress
	_ = callLifecycleHook(ctx, apiURL, event.SyncJobID, "events", map[string]interface{}{
		"type":    "info",
		"message": fmt.Sprintf("Starting collection for account %d in region %s", event.CloudAccountID, event.Region),
		"metadata": map[string]interface{}{
			"provider": "aws",
			"region":   event.Region,
		},
	})

	ingestURL := fmt.Sprintf("%s/api/v1/caasm/ingest", apiURL)
	req, err := http.NewRequestWithContext(ctx, "POST", ingestURL, bytes.NewBuffer(payloadBytes))
	if err != nil {
		return "", fmt.Errorf("failed to create request: %v", err)
	}
	req.Header.Set("Content-Type", "application/json")
	// TODO: Add Auth Header

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("failed to send data to backend: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("backend returned status: %d", resp.StatusCode)
	}

	return fmt.Sprintf("Successfully ingested %d assets", len(result.Assets)), nil
}

func callLifecycleHook(ctx context.Context, apiURL string, jobID uint64, action string, payload interface{}) error {
	url := fmt.Sprintf("%s/api/v1/caasm/sync-jobs/%d/%s", apiURL, jobID, action)

	var body []byte
	var err error
	if payload != nil {
		body, err = json.Marshal(payload)
		if err != nil {
			return err
		}
	}

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

	client := &http.Client{Timeout: 5 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("status code %d", resp.StatusCode)
	}
	return nil
}

func main() {
	// NOTE: This is a standalone AWS Lambda handler entrypoint. It can be built and
	// deployed independently, but the main API service does not currently invoke it.
	// The in-product "serverless" sync mode still goes through invokeLambdaSimulated(...)
	// in internal/modules/caasm/service/service.go.
	lambda.Start(HandleRequest)
}