#!/usr/bin/env python3
"""
Vulnerability Scan and Remediation Agent

This script runs a Trivy security scan on the repository, identifies vulnerabilities
above a severity threshold, and uses OpenHands agents to create PRs with fixes.

Usage:
    python scan_and_remediate.py --scan-only    # Only run Trivy scan, output results
    python scan_and_remediate.py --remediate    # Run remediation on existing scan results
    python scan_and_remediate.py                # Full scan + remediation (legacy mode)

Environment Variables:
    LLM_API_KEY: API key for the LLM (required for remediation)
    LLM_MODEL: Language model to use (default: anthropic/claude-sonnet-4-5-20250929)
    LLM_BASE_URL: Optional base URL for LLM API
    GITHUB_TOKEN: GitHub token for API access and creating PRs (required for remediation)
    REPO_NAME: Repository name in format owner/repo (required for remediation)
    SEVERITY_THRESHOLD: Minimum severity to remediate (default: HIGH)
    MAX_VULNERABILITIES: Maximum vulnerabilities to fix per run (default: 5, 0=unlimited)
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

# Set up basic logging for scan-only mode (no openhands dependency)
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
basic_logger = logging.getLogger(__name__)

SEVERITY_ORDER = ["CRITICAL", "HIGH", "MEDIUM", "LOW"]


@dataclass
class Vulnerability:
    """Represents a security vulnerability."""

    vuln_id: str
    package_name: str
    installed_version: str
    fixed_version: str | None
    severity: str
    title: str
    description: str
    target: str

    @property
    def has_fix(self) -> bool:
        return self.fixed_version is not None and self.fixed_version != ""


def get_required_env(name: str) -> str:
    """Get a required environment variable."""
    value = os.getenv(name)
    if not value:
        raise ValueError(f"{name} environment variable is required")
    return value


def run_trivy_scan(repo_path: str, logger) -> dict:
    """Run Trivy scan on the repository and return results."""
    logger.info("Running Trivy security scan...")

    output_file = Path(repo_path) / "trivy-results.json"

    cmd = [
        "trivy",
        "fs",
        "--format",
        "json",
        "--output",
        str(output_file),
        "--severity",
        "CRITICAL,HIGH,MEDIUM,LOW",
        repo_path,
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
        if result.returncode != 0 and not output_file.exists():
            logger.error(f"Trivy scan failed: {result.stderr}")
            raise RuntimeError(f"Trivy scan failed: {result.stderr}")

        with open(output_file) as f:
            return json.load(f)

    except subprocess.TimeoutExpired:
        raise RuntimeError("Trivy scan timed out after 10 minutes")
    except json.JSONDecodeError as e:
        raise RuntimeError(f"Failed to parse Trivy output: {e}")


def parse_vulnerabilities(trivy_output: dict) -> list[Vulnerability]:
    """Parse Trivy output into Vulnerability objects."""
    vulnerabilities = []

    results = trivy_output.get("Results", [])
    for result in results:
        target = result.get("Target", "unknown")
        vulns = result.get("Vulnerabilities", [])

        for vuln in vulns:
            vulnerabilities.append(
                Vulnerability(
                    vuln_id=vuln.get("VulnerabilityID", ""),
                    package_name=vuln.get("PkgName", ""),
                    installed_version=vuln.get("InstalledVersion", ""),
                    fixed_version=vuln.get("FixedVersion"),
                    severity=vuln.get("Severity", "UNKNOWN"),
                    title=vuln.get("Title", ""),
                    description=vuln.get("Description", ""),
                    target=target,
                )
            )

    return vulnerabilities


def filter_vulnerabilities(
    vulnerabilities: list[Vulnerability],
    severity_threshold: str,
    max_count: int,
) -> list[Vulnerability]:
    """Filter vulnerabilities by severity threshold and limit count.
    
    Note: severity_threshold should be validated before calling this function.
    """
    # Ensure valid threshold (defensive, validation should happen earlier)
    if severity_threshold not in SEVERITY_ORDER:
        severity_threshold = "HIGH"

    threshold_index = SEVERITY_ORDER.index(severity_threshold)
    allowed_severities = set(SEVERITY_ORDER[: threshold_index + 1])

    # Filter by severity and only include those with fixes
    filtered = [
        v
        for v in vulnerabilities
        if v.severity in allowed_severities and v.has_fix
    ]

    # Sort by severity (most critical first)
    filtered.sort(key=lambda v: SEVERITY_ORDER.index(v.severity))

    # Limit count
    if max_count > 0:
        filtered = filtered[:max_count]

    return filtered


def create_remediation_prompt(vuln: Vulnerability, repo_name: str) -> str:
    """Create the prompt for the remediation agent."""
    return f"""You are a security engineer tasked with fixing a vulnerability in a repository.

## Repository
{repo_name}

## Vulnerability Details
- **ID**: {vuln.vuln_id}
- **Package**: {vuln.package_name}
- **Severity**: {vuln.severity}
- **Current Version**: {vuln.installed_version}
- **Fixed Version**: {vuln.fixed_version}
- **Location**: {vuln.target}
- **Title**: {vuln.title}
- **Description**: {vuln.description}

## Your Task

1. **Analyze** the vulnerability and understand what needs to be fixed

2. **Find ALL dependency files across the repository**
   **CRITICAL**: Search the repository for ALL dependency files that contain the vulnerable package. Do NOT assume dependencies only exist in the root directory.

   ```bash
   # Find ALL dependency files containing the package (excludes dependency/build directories)
   find . \( -name node_modules -o -name .venv -o -name venv -o -name vendor -o -name .git -o -name __pycache__ -o -name dist -o -name build \) -prune -o -name "pyproject.toml" -exec grep -l "{vuln.package_name}" {} + 2>/dev/null
   find . \( -name node_modules -o -name .venv -o -name venv -o -name vendor -o -name .git \) -prune -o -name "requirements*.txt" -exec grep -l "{vuln.package_name}" {} + 2>/dev/null
   find . \( -name node_modules -o -name .venv -o -name venv -o -name vendor -o -name .git \) -prune -o -name "package.json" -exec grep -l "{vuln.package_name}" {} + 2>/dev/null
   ```

   Update {vuln.package_name} from {vuln.installed_version} to {vuln.fixed_version} in **EVERY** file found.

3. **CRITICAL: Sync/regenerate ALL lockfiles in the repository**
   After updating the version in ALL manifest files, you MUST regenerate ALL corresponding lockfiles. Do NOT manually edit lockfiles.

   **CRITICAL**: First, find ALL lockfiles (excluding dependency/build directories):
   ```bash
   find . \\( -name node_modules -o -name .venv -o -name venv -o -name vendor -o -name .git \\) -prune -o \\( -name "poetry.lock" -o -name "uv.lock" -o -name "package-lock.json" -o -name "yarn.lock" -o -name "pnpm-lock.yaml" -o -name "Cargo.lock" -o -name "go.sum" \\) -print
   ```

   You MUST regenerate EVERY lockfile found, not just those in the root directory.

   **IMPORTANT**: To avoid unnecessary diff noise, you MUST detect and use the same tool version that originally generated each lockfile. Each lockfile contains a version header that indicates which tool version was used.

   **Note**: This agent runs in an isolated execution environment (container), so installing specific tool versions with `--force` will not affect other projects or system-wide installations.

   For EACH lockfile found, `cd` to its directory and regenerate it:

   - **Poetry (pyproject.toml + poetry.lock)**:
     1. `cd` to the directory containing the lockfile
     2. Extract version: `grep -m1 "^# This file is automatically @generated by Poetry" poetry.lock | sed 's/.*Poetry \\([0-9.]*\\).*/\\1/'`
     3. If a version is found, install it: `pipx install poetry==$POETRY_VERSION --force`
     4. Verify installation: `poetry --version | grep "$POETRY_VERSION"` (proceed only if successful)
     5. If version extraction fails or returns empty, proceed with the currently installed version and note this in your output
     6. Run: `poetry lock --no-update` or `poetry update {vuln.package_name}`

   - **uv (pyproject.toml + uv.lock)**:
     1. `cd` to the directory containing the lockfile
     2. Extract version: `grep -m1 "^# This file was autogenerated by uv" uv.lock | sed 's/.*uv version \\([0-9.]*\\).*/\\1/'`
     3. If a version is found, install it: `pipx install uv==$UV_VERSION --force`
     4. Verify installation: `uv --version | grep "$UV_VERSION"` (proceed only if successful)
     5. If version extraction fails or returns empty, proceed with the currently installed version and note this in your output
     6. Run: `uv lock --upgrade-package {vuln.package_name}` or `uv sync`

   - **npm (package.json + package-lock.json)**: `cd` to directory, run `npm install` or `npm update {vuln.package_name}`
   - **yarn (package.json + yarn.lock)**: `cd` to directory, run `yarn install` or `yarn upgrade {vuln.package_name}`
   - **pnpm (package.json + pnpm-lock.yaml)**: `cd` to directory, run `pnpm install` or `pnpm update {vuln.package_name}`
   - **pip (requirements.txt)**: Update the version directly in requirements.txt
   - **Go (go.mod + go.sum)**: `cd` to directory, run `go mod tidy`
   - **Cargo (Cargo.toml + Cargo.lock)**: `cd` to directory, run `cargo update -p {vuln.package_name}`
   - **Maven (pom.xml)**: Update the version directly in pom.xml
   - **Gradle**: Update the version in build.gradle/build.gradle.kts

4. **Verify** the change doesn't break the build (run any available build/test commands)
5. **Create a branch** named `fix/{vuln.vuln_id.lower()}`
6. **Commit** your changes with a clear message explaining the security fix
7. **Push** the branch to origin
8. **Create a Pull Request** using the GitHub CLI:
   ```bash
   gh pr create --title "fix: {vuln.vuln_id} - Update {vuln.package_name} to {vuln.fixed_version}" \\
     --body "## Security Fix

This PR addresses {vuln.vuln_id} ({vuln.severity} severity).

### Vulnerability
{vuln.title}

### Changes
- Updated `{vuln.package_name}` from `{vuln.installed_version}` to `{vuln.fixed_version}`

### References
- https://nvd.nist.gov/vuln/detail/{vuln.vuln_id}
"
   ```

## Important Notes
- Do NOT modify any code beyond what's necessary for the fix
- Do NOT manually edit lockfiles - always use the package manager commands above
- If the package update requires other dependency changes, include them
- If you encounter conflicts or issues, document them in the PR description
- Always test that the fix doesn't break the build before creating the PR
"""


def create_agent(config: dict, logger):
    """Create and configure the remediation agent."""
    from openhands.sdk import LLM, Agent
    from openhands.tools.preset.default import get_default_condenser, get_default_tools

    llm_config = {
        "model": config["model"],
        "api_key": config["api_key"],
        "usage_id": "vulnerability_remediation",
        "drop_params": True,
    }
    if config.get("base_url"):
        llm_config["base_url"] = config["base_url"]

    llm = LLM(**llm_config)

    return Agent(
        llm=llm,
        tools=get_default_tools(enable_browser=False),
        system_prompt_kwargs={"cli_mode": True},
        condenser=get_default_condenser(
            llm=llm.model_copy(update={"usage_id": "condenser"})
        ),
    )


def remediate_vulnerability(
    agent,
    vuln: Vulnerability,
    repo_name: str,
    secrets: dict[str, str],
    logger,
) -> dict:
    """Run the remediation agent for a single vulnerability."""
    from openhands.sdk import Conversation

    logger.info(f"Remediating {vuln.vuln_id} ({vuln.severity}): {vuln.package_name}")

    prompt = create_remediation_prompt(vuln, repo_name)
    cwd = os.getcwd()

    conversation = Conversation(
        agent=agent,
        workspace=cwd,
        secrets=secrets,
    )

    try:
        conversation.send_message(prompt)
        conversation.run()

        metrics = conversation.conversation_stats.get_combined_metrics()

        return {
            "vuln_id": vuln.vuln_id,
            "package": vuln.package_name,
            "severity": vuln.severity,
            "status": "completed",
            "cost": metrics.accumulated_cost,
        }

    except Exception as e:
        logger.error(f"Failed to remediate {vuln.vuln_id}: {type(e).__name__}: {e}")
        return {
            "vuln_id": vuln.vuln_id,
            "package": vuln.package_name,
            "severity": vuln.severity,
            "status": "failed",
            "error": str(e),
            "error_type": type(e).__name__,
        }


def save_report(vulnerabilities: list[Vulnerability], results: list[dict], logger) -> None:
    """Save the remediation report."""
    report = {
        "total_vulnerabilities_found": len(vulnerabilities),
        "remediation_results": results,
        "summary": {
            "completed": len([r for r in results if r["status"] == "completed"]),
            "failed": len([r for r in results if r["status"] == "failed"]),
            "total_cost": sum(r.get("cost", 0) for r in results),
        },
    }

    with open("remediation-report.json", "w") as f:
        json.dump(report, f, indent=2)

    logger.info("Remediation report saved to remediation-report.json")


def save_scan_results(
    all_vulns: list[Vulnerability],
    vulns_to_fix: list[Vulnerability],
    severity_threshold: str,
    logger,
) -> None:
    """Save scan results for the action to read."""
    scan_results = {
        "total_vulnerabilities": len(all_vulns),
        "vulnerabilities_to_fix": len(vulns_to_fix),
        "severity_threshold": severity_threshold,
        "vulnerabilities": [
            {
                "vuln_id": v.vuln_id,
                "package_name": v.package_name,
                "installed_version": v.installed_version,
                "fixed_version": v.fixed_version,
                "severity": v.severity,
                "title": v.title,
                "target": v.target,
            }
            for v in vulns_to_fix
        ],
    }

    with open("scan-results.json", "w") as f:
        json.dump(scan_results, f, indent=2)

    logger.info("Scan results saved to scan-results.json")


def validate_severity_threshold(severity_threshold: str, logger) -> str:
    """Validate severity threshold, warn and default to HIGH if invalid."""
    if severity_threshold not in SEVERITY_ORDER:
        logger.warning(
            f"Invalid severity threshold '{severity_threshold}'. "
            f"Valid options: {', '.join(SEVERITY_ORDER)}. Defaulting to 'HIGH'."
        )
        return "HIGH"
    return severity_threshold


def run_scan_only():
    """Run only the Trivy scan and save results (no OpenHands dependency)."""
    logger = basic_logger
    logger.info("Running vulnerability scan (scan-only mode)...")

    severity_threshold = os.getenv("SEVERITY_THRESHOLD", "HIGH")
    severity_threshold = validate_severity_threshold(severity_threshold, logger)
    max_vulns = int(os.getenv("MAX_VULNERABILITIES", "5"))

    logger.info(f"Severity threshold: {severity_threshold}")
    logger.info(f"Max vulnerabilities: {max_vulns if max_vulns > 0 else 'unlimited'}")

    # Run Trivy scan
    repo_path = os.getcwd()
    trivy_output = run_trivy_scan(repo_path, logger)

    # Parse and filter vulnerabilities
    all_vulns = parse_vulnerabilities(trivy_output)
    logger.info(f"Found {len(all_vulns)} total vulnerabilities")

    vulns_to_fix = filter_vulnerabilities(all_vulns, severity_threshold, max_vulns)
    logger.info(
        f"Filtered to {len(vulns_to_fix)} vulnerabilities "
        f"({severity_threshold}+ with available fixes)"
    )

    # Save scan results for the action to read
    save_scan_results(all_vulns, vulns_to_fix, severity_threshold, logger)

    print(f"\n=== Vulnerability Scan Summary ===")
    print(f"Total Vulnerabilities Found: {len(all_vulns)}")
    print(f"Vulnerabilities to Remediate: {len(vulns_to_fix)}")

    if not vulns_to_fix:
        logger.info("✅ No vulnerabilities found that need remediation!")
    else:
        logger.info(f"🔍 Found {len(vulns_to_fix)} vulnerabilities to remediate")


def run_remediation():
    """Run remediation using existing scan results (requires OpenHands)."""
    from openhands.sdk import LLM, Agent, Conversation, get_logger
    from openhands.tools.preset.default import get_default_condenser, get_default_tools

    logger = get_logger(__name__)
    logger.info("Running vulnerability remediation...")

    # Load scan results
    scan_results_file = Path("scan-results.json")
    if not scan_results_file.exists():
        logger.error("No scan results found. Run with --scan-only first.")
        sys.exit(1)

    with open(scan_results_file) as f:
        scan_results = json.load(f)

    vulns_data = scan_results.get("vulnerabilities", [])
    if not vulns_data:
        logger.info("No vulnerabilities to remediate!")
        return

    # Convert to Vulnerability objects
    vulns_to_fix = [
        Vulnerability(
            vuln_id=v["vuln_id"],
            package_name=v["package_name"],
            installed_version=v["installed_version"],
            fixed_version=v["fixed_version"],
            severity=v["severity"],
            title=v["title"],
            description="",  # Not stored in scan results
            target=v["target"],
        )
        for v in vulns_data
    ]

    # Get configuration
    config = {
        "api_key": get_required_env("LLM_API_KEY"),
        "model": os.getenv("LLM_MODEL", "anthropic/claude-sonnet-4-5-20250929"),
        "base_url": os.getenv("LLM_BASE_URL"),
    }

    github_token = get_required_env("GITHUB_TOKEN")
    repo_name = get_required_env("REPO_NAME")

    logger.info(f"Repository: {repo_name}")
    logger.info(f"Remediating {len(vulns_to_fix)} vulnerabilities...")

    # Create agent and remediate
    agent = create_agent(config, logger)
    secrets = {
        "LLM_API_KEY": config["api_key"],
        "GITHUB_TOKEN": github_token,
    }

    results = []
    for vuln in vulns_to_fix:
        result = remediate_vulnerability(agent, vuln, repo_name, secrets, logger)
        results.append(result)

    # Save report
    save_report(vulns_to_fix, results, logger)

    # Print summary
    completed = len([r for r in results if r["status"] == "completed"])
    failed = len([r for r in results if r["status"] == "failed"])
    total_cost = sum(r.get("cost", 0) for r in results)

    print("\n=== Vulnerability Remediation Summary ===")
    print(f"Attempted Remediations: {len(results)}")
    print(f"Completed: {completed}")
    print(f"Failed: {failed}")
    print(f"Total Cost: ${total_cost:.6f}")

    if failed > 0:
        logger.warning(f"{failed} remediations failed - check logs for details")
        sys.exit(1)

    logger.info("Vulnerability remediation completed successfully")


def main():
    """Run vulnerability scan and remediation (legacy mode - full workflow).
    
    This combines scan-only and remediate modes for backward compatibility.
    Prefer using --scan-only and --remediate separately for better control.
    """
    # Run scan first
    run_scan_only()
    
    # Check if there are vulnerabilities to fix
    scan_results_file = Path("scan-results.json")
    if scan_results_file.exists():
        with open(scan_results_file) as f:
            scan_results = json.load(f)
        if scan_results.get("vulnerabilities_to_fix", 0) > 0:
            # Run remediation
            run_remediation()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Vulnerability Scan and Remediation Agent"
    )
    parser.add_argument(
        "--scan-only",
        action="store_true",
        help="Only run Trivy scan, save results without starting agent",
    )
    parser.add_argument(
        "--remediate",
        action="store_true",
        help="Run remediation using existing scan results",
    )

    args = parser.parse_args()

    if args.scan_only:
        run_scan_only()
    elif args.remediate:
        run_remediation()
    else:
        # Legacy mode: full scan + remediation
        main()