src/iac_secfix_agents/agents.py

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .config import RuntimeConfig
from .llm import OpenAICompatibleLLM, make_pydantic_agent, make_pydantic_openai_model, parse_fix_candidate
from .policy import build_rule_hints_md, enrich_findings
from .schemas import (
    ApprovalDecision,
    ComplianceReport,
    EnrichedFindings,
    Finding,
    FixCandidate,
    IngestRequest,
    ScanBundle,
    ValidationResult,
)


@dataclass(slots=True)
class PydanticAgentBundle:
    ingest: Any | None = None
    policy_context: Any | None = None
    patch: Any | None = None
    validation: Any | None = None
    compliance: Any | None = None
    approval: Any | None = None

    @property
    def available(self) -> bool:
        return any(
            agent is not None
            for agent in [
                self.ingest,
                self.policy_context,
                self.patch,
                self.validation,
                self.compliance,
                self.approval,
            ]
        )


def create_pydantic_agent_bundle(config: RuntimeConfig) -> PydanticAgentBundle:
    """Build PydanticAI agents when the installed API supports it."""

    model = make_pydantic_openai_model(config)
    return PydanticAgentBundle(
        ingest=make_pydantic_agent(
            model,
            ScanBundle,
            "Identify file type and return scanner bundle. Deterministic scanner tools supply findings.",
        ),
        policy_context=make_pydantic_agent(
            model,
            EnrichedFindings,
            "Map scanner findings to policy descriptions. Do not invent policy IDs.",
        ),
        patch=make_pydantic_agent(model, FixCandidate, "Return only a valid FixCandidate JSON object."),
        validation=make_pydantic_agent(model, ValidationResult, "Summarize deterministic validation results only."),
        compliance=make_pydantic_agent(model, ComplianceReport, "Compute compliance summary only from scan summaries."),
        approval=make_pydantic_agent(
            model,
            ApprovalDecision,
            "Return needs_human_review unless deterministic validation says accept and auto approval is enabled.",
        ),
    )


def shard_for_context(file_content: str, findings: list[Finding], max_chars: int) -> str:
    if len(file_content) <= max_chars:
        return file_content

    lines = file_content.splitlines()
    selected: list[str] = ["# Non-targeted blocks omitted for context budget."]
    used_ranges: list[tuple[int, int]] = []

    for finding in findings:
        start = max((finding.line_start or 1) - 20, 1)
        end = min((finding.line_end or start) + 20, len(lines))
        if any(not (end < old_start or start > old_end) for old_start, old_end in used_ranges):
            continue
        used_ranges.append((start, end))
        selected.append(f"\n# Context lines {start}-{end}")
        selected.extend(lines[start - 1 : end])
        if len("\n".join(selected)) >= max_chars:
            break

    return "\n".join(selected)[:max_chars]


def fix_candidate_schema_text() -> str:
    return """
Return exactly one JSON object:
{
  "patch": "unified diff string",
  "fixed_file": "complete corrected file content",
  "resolved_policy_ids": ["CKV_..."],
  "explanation": "short summary",
  "verification_commands": ["checkov -f <file> -o json"],
  "risk_notes": ["human-review note if needed"],
  "requires_human_approval": true
}
No markdown fences. No prose outside JSON.
""".strip()


def build_patch_prompt(
    enriched: EnrichedFindings,
    retry_feedback: str | None = None,
    max_context_chars: int = 18000,
    max_finding_records: int = 24,
) -> str:
    iac_type = enriched.file.iac_type
    lang = "hcl" if iac_type == "terraform" else "yaml"
    findings_for_json = enriched.findings[:max_finding_records]
    omitted = max(0, len(enriched.findings) - len(findings_for_json))

    scanner_report_json = json.dumps(
        [finding.model_dump(mode="json") for finding in findings_for_json],
        ensure_ascii=False,
        indent=2,
    )
    rule_ids = [finding.rule_id for finding in enriched.findings]
    original_context = shard_for_context(enriched.file.file_content, list(enriched.findings), max_context_chars)

    prompt = f"""
Original file type: {iac_type}
Original file language: {lang}

<original_file>
{original_context}
</original_file>

<targeted_rule_ids>
{json.dumps(sorted(set(rule_ids)))}
</targeted_rule_ids>

<scanner_findings_json>
{scanner_report_json}
</scanner_findings_json>

<scanner_findings_omitted_due_context_budget>
{omitted}
</scanner_findings_omitted_due_context_budget>

<targeted_rules_with_fix_templates>
{build_rule_hints_md(rule_ids)}
</targeted_rules_with_fix_templates>

Task:
Create the smallest safe in-file remediation that resolves as many targeted rule_id values as possible.

Hard requirements:
- Start from the ORIGINAL file. Add or modify only attributes directly needed by the matching fix templates.
- fixed_file must contain the complete corrected file, including unchanged lines.
- Do not delete, rename, or reorder resources, modules, providers, containers, or attributes unless the rule requires it.
- Do not invent placeholder ARNs, account IDs, bucket names, topics, roles, regions, KMS keys, CIDRs, or image tags.
- If a rule requires an external dependency or unknown business decision, leave it unresolved and set requires_human_approval=true.
- Do not introduce CRITICAL or HIGH findings.
- Do not add checkov:skip, nosec, tfsec:ignore, or scanner suppression comments.
- Before returning, list only rule IDs that your fixed_file actually resolves in resolved_policy_ids.

{fix_candidate_schema_text()}
""".strip()

    if retry_feedback:
        prompt += f"\n\n<previous_validation_feedback>\n{retry_feedback[:4000]}\n</previous_validation_feedback>"

    return prompt


def compress_feedback(result: ValidationResult, attempt: int, max_retries: int) -> str:
    payload = {
        "attempt": attempt + 1,
        "max_retries": max_retries,
        "retry_reason": result.retry_reason,
        "targeted_still_unresolved": result.targeted_unresolved,
        "new_critical_high_count": result.new_critical_high_count,
        "changed_lines_ratio": round(result.changed_lines_ratio, 3),
        "post_scan_summary": result.post_summary,
    }
    feedback = json.dumps(payload, ensure_ascii=False)
    if result.targeted_unresolved:
        feedback += "\n\nUnresolved-rule fix templates:\n"
        feedback += build_rule_hints_md(result.targeted_unresolved)
    feedback += (
        "\n\nRetry instructions:\n"
        "- Start again from the original file, not the previous attempt.\n"
        "- Apply the smallest safe edit for the unresolved rules.\n"
        "- Remove rule IDs from resolved_policy_ids when validation did not resolve them.\n"
    )
    return feedback[:4000]


class IngestAgent:
    def run(self, request: IngestRequest, tmp_dir: Path | None = None) -> ScanBundle:
        from .scanners import unified_scan

        return unified_scan(request, tmp_dir=tmp_dir)


class PolicyContextAgent:
    def run(self, bundle: ScanBundle) -> EnrichedFindings:
        return enrich_findings(bundle)


class PatchAgent:
    def __init__(self, config: RuntimeConfig, log_dir: Path):
        self.llm = OpenAICompatibleLLM(config)
        self.log_dir = log_dir
        self.log_dir.mkdir(parents=True, exist_ok=True)

    def run(
        self,
        enriched: EnrichedFindings,
        retry_feedback: str | None = None,
        max_context_chars: int = 18000,
    ) -> tuple[FixCandidate | None, bool, str]:
        prompt = build_patch_prompt(enriched, retry_feedback, max_context_chars=max_context_chars)
        messages = [
            {
                "role": "system",
                "content": (
                    "You are IaC-SecFix. Return only valid JSON matching the FixCandidate schema. "
                    "Do not include markdown, prose, analysis, or code fences."
                ),
            },
            {"role": "user", "content": prompt},
        ]
        result = self.llm.chat(messages, json_mode=True)
        self.log_dir.joinpath("last_patch_raw.txt").write_text(result.text[:20000], errors="ignore")
        candidate, schema_valid = parse_fix_candidate(result.text)
        return candidate, schema_valid, result.text


class ValidationAgent:
    def run(
        self,
        request: IngestRequest,
        candidate: FixCandidate,
        targeted: list[str],
        tmp_dir: Path | None,
        max_changed_lines_ratio: float,
    ) -> ValidationResult:
        from .scanners import validate_candidate

        return validate_candidate(
            request,
            candidate,
            targeted,
            tmp_dir=tmp_dir,
            max_changed_lines_ratio=max_changed_lines_ratio,
        )


class ComplianceAgent:
    def run(self, pre_summary: dict[str, object], result: ValidationResult) -> ComplianceReport:
        before_failed = max(int(pre_summary.get("failed", 0) or 0), len(result.targeted_resolved) + len(result.targeted_unresolved), 1)
        after_failed = int(result.post_summary.get("failed", 0) or 0)
        targeted_total = max(len(result.targeted_resolved) + len(result.targeted_unresolved), 1)
        return ComplianceReport(
            cis_pass_rate_before=max(0.0, 1.0 - before_failed / max(before_failed, 1)),
            cis_pass_rate_after=max(0.0, 1.0 - after_failed / max(before_failed, 1)),
            network_exposure_reduction=1.0 if result.targeted_resolved else 0.0,
            iam_privilege_reduction=1.0 if result.targeted_resolved else 0.0,
            targeted_resolution_rate=len(result.targeted_resolved) / targeted_total,
            scanner_clean=after_failed == 0,
        )


class HumanApprovalGate:
    def __init__(self, auto_approve_validated: bool = False):
        self.auto_approve_validated = auto_approve_validated

    def run(
        self,
        patch: FixCandidate,
        result: ValidationResult,
        compliance: ComplianceReport,
        n_retries: int,
    ) -> ApprovalDecision:
        can_auto_approve = (
            self.auto_approve_validated
            and result.verdict == "accept"
            and not patch.requires_human_approval
        )
        if can_auto_approve:
            return ApprovalDecision(
                action="approved",
                reason="Deterministic validation accepted the patch.",
                patch=patch,
                validation=result,
                compliance=compliance,
                n_retries=n_retries,
            )
        return ApprovalDecision(
            action="needs_human_review",
            reason="Validated candidate requires human approval before export.",
            patch=patch,
            validation=result,
            compliance=compliance,
            n_retries=n_retries,
        )