package report

import (
	"context"
	"fmt"
	"time"

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/llm"
	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/pipeline"
	"github.com/google/uuid"
)

// ReportAgent generates professional pentest reports using LLM.
type ReportAgent struct {
	provider llm.Provider
}

// NewReportAgent creates a new report agent.
func NewReportAgent(provider llm.Provider) *ReportAgent {
	return &ReportAgent{provider: provider}
}

// Generate produces a full PentestReport from campaign data.
func (r *ReportAgent) Generate(ctx context.Context, campaign pipeline.Campaign, findings []pipeline.ClassifiedFinding, plan *pipeline.AttackPlan, results []pipeline.ExecutionResult) (*pipeline.PentestReport, error) {
	report := &pipeline.PentestReport{
		ID:          uuid.New(),
		CampaignID:  campaign.ID,
		Target:      campaign.Target,
		Objective:   campaign.Objective,
		GeneratedAt: time.Now(),
	}

	// Generate executive summary
	execSummary, err := r.generateSection(ctx, "executive_summary", campaign, findings)
	if err == nil {
		report.ExecutiveSummary = execSummary
	}

	// Generate finding writeups
	for _, f := range findings {
		reportFinding := pipeline.ReportFinding{
			ID:                 f.ID,
			Title:              f.Title,
			Severity:           f.Severity,
			CVSSScore:          f.CVSSScore,
			CVSSVector:         f.CVSSVector,
			Description:        f.Description,
			Evidence:           f.Evidence,
			AffectedComponents: []string{f.Target},
		}

		// Generate remediation for each finding
		remediation, err := r.generateRemediation(ctx, f)
		if err == nil {
			reportFinding.Remediation = remediation
		}

		report.Findings = append(report.Findings, reportFinding)
	}

	// Generate attack narrative
	if plan != nil {
		narrative, err := r.generateNarrative(ctx, campaign, plan, results)
		if err == nil {
			report.AttackNarrative = narrative
		}
	}

	// Build risk summary
	report.RiskSummary = buildRiskSummary(findings)

	// Generate remediation plan
	report.RemediationPlan = buildRemediationPlan(findings)

	return report, nil
}

func (r *ReportAgent) generateSection(ctx context.Context, section string, campaign pipeline.Campaign, findings []pipeline.ClassifiedFinding) (string, error) {
	var prompt string
	switch section {
	case "executive_summary":
		prompt = fmt.Sprintf(
			"Write a 2-3 paragraph executive summary for a penetration test report.\nTarget: %s\nObjective: %s\nTotal findings: %d critical, high-severity findings. Focus on business risk, not technical details. Write for a non-technical audience.",
			campaign.Target, campaign.Objective, len(findings),
		)
	default:
		return "", fmt.Errorf("unknown section: %s", section)
	}

	resp, err := r.provider.Complete(ctx, llm.CompletionRequest{
		SystemPrompt: "You are a professional penetration testing report writer. Write clear, concise, and actionable content.",
		Messages:     []llm.Message{{Role: "user", Content: prompt}},
		MaxTokens:    2048,
		Temperature:  0.3,
	})
	if err != nil {
		return "", err
	}

	return resp.Content, nil
}

func (r *ReportAgent) generateRemediation(ctx context.Context, finding pipeline.ClassifiedFinding) (string, error) {
	resp, err := r.provider.Complete(ctx, llm.CompletionRequest{
		SystemPrompt: "You are a security remediation expert. Provide specific, actionable remediation steps.",
		Messages: []llm.Message{
			{
				Role:    "user",
				Content: fmt.Sprintf("Provide remediation steps for this vulnerability:\nTitle: %s\nSeverity: %s\nCVSS: %.1f\nDescription: %s", finding.Title, finding.Severity, finding.CVSSScore, finding.Description),
			},
		},
		MaxTokens:   1024,
		Temperature: 0.2,
	})
	if err != nil {
		return "", err
	}
	return resp.Content, nil
}

func (r *ReportAgent) generateNarrative(ctx context.Context, campaign pipeline.Campaign, plan *pipeline.AttackPlan, results []pipeline.ExecutionResult) (string, error) {
	resp, err := r.provider.Complete(ctx, llm.CompletionRequest{
		SystemPrompt: "You are a penetration testing report writer. Write the attack narrative as a sequence of events, telling the story of this penetration test from initial recon to final findings.",
		Messages: []llm.Message{
			{
				Role:    "user",
				Content: fmt.Sprintf("Write an attack narrative for target %s. %d attack paths were tested with %d execution results. Reasoning: %s", campaign.Target, len(plan.Paths), len(results), plan.Reasoning),
			},
		},
		MaxTokens:   2048,
		Temperature: 0.3,
	})
	if err != nil {
		return "", err
	}
	return resp.Content, nil
}

func buildRiskSummary(findings []pipeline.ClassifiedFinding) pipeline.RiskSummary {
	summary := pipeline.RiskSummary{}
	for _, f := range findings {
		switch f.Severity {
		case pipeline.SeverityCritical:
			summary.CriticalCount++
		case pipeline.SeverityHigh:
			summary.HighCount++
		case pipeline.SeverityMedium:
			summary.MediumCount++
		case pipeline.SeverityLow:
			summary.LowCount++
		case pipeline.SeverityInformational:
			summary.InfoCount++
		}
	}

	switch {
	case summary.CriticalCount > 0:
		summary.OverallRisk = "critical"
	case summary.HighCount > 0:
		summary.OverallRisk = "high"
	case summary.MediumCount > 0:
		summary.OverallRisk = "medium"
	default:
		summary.OverallRisk = "low"
	}

	return summary
}

func buildRemediationPlan(findings []pipeline.ClassifiedFinding) []pipeline.RemediationItem {
	var items []pipeline.RemediationItem
	for i, f := range findings {
		items = append(items, pipeline.RemediationItem{
			Priority: i + 1,
			Finding:  f.Title,
			Action:   "Remediate " + f.Title,
			Effort:   estimateEffort(f.Severity),
			Impact:   string(f.Severity),
		})
	}
	return items
}

func estimateEffort(severity pipeline.Severity) string {
	switch severity {
	case pipeline.SeverityCritical:
		return "immediate"
	case pipeline.SeverityHigh:
		return "1-2 days"
	case pipeline.SeverityMedium:
		return "1 week"
	default:
		return "low priority"
	}
}
