"""VulnClaw PoC Builder — auto-generate runnable PoC scripts for findings."""

from __future__ import annotations

import re
from datetime import datetime
from pathlib import Path

from vulnclaw.agent.context import SessionState, VulnerabilityFinding
from vulnclaw.report.verifier import PoCGenerator

PYTHON_POC_TEMPLATE = '''\
#!/usr/bin/env python3
"""
PoC: {{ title }}
Severity: {{ severity }}
{% if cve %}CVE: {{ cve }}
{% endif %}Vulnerability Type: {{ vuln_type or "generic" }}
Target: {{ target }}
Generated by VulnClaw on {{ generated_at }}

Notes:
- This PoC is generated from the report finding and uses vulnerability-specific verification logic when available.
- Review the payload/parameters before running against real systems.
"""

{{ poc_body }}
'''


_VULN_TYPE_ALIASES: dict[str, str] = {
    "sqli": "sql_injection",
    "sql injection": "sql_injection",
    "sql_injection": "sql_injection",
    "xss": "xss",
    "cross_site_scripting": "xss",
    "cross-site_scripting": "xss",
    "cross site scripting": "xss",
    "rce": "command_injection",
    "command_injection": "command_injection",
    "command injection": "command_injection",
    "cmd_injection": "command_injection",
    "lfi": "lfi",
    "local_file_inclusion": "lfi",
    "local file inclusion": "lfi",
    "info leak": "info_disclosure",
    "information_disclosure": "info_disclosure",
    "information disclosure": "info_disclosure",
    "debug": "debug_mode",
    "debug_mode": "debug_mode",
    "sensitive_file": "sensitive_file",
    "sensitive file": "sensitive_file",
}


def _sanitize_filename_component(value: str, max_length: int = 40) -> str:
    sanitized = re.sub(r'[<>:"/\\|?*]+', "_", value or "")
    sanitized = re.sub(r"\s+", "_", sanitized).strip(" ._")
    sanitized = re.sub(r"_+", "_", sanitized)
    if not sanitized:
        sanitized = "finding"
    return sanitized[:max_length]


def _normalize_vuln_type(vuln_type: str) -> str:
    normalized = (vuln_type or "").strip().lower().replace("-", " ").replace("/", " ")
    compact = re.sub(r"\s+", " ", normalized)
    underscored = compact.replace(" ", "_")
    return _VULN_TYPE_ALIASES.get(compact, _VULN_TYPE_ALIASES.get(underscored, underscored))


def _extract_target(finding: VulnerabilityFinding, fallback_target: str) -> str:
    for field in (finding.evidence, finding.description):
        if not field:
            continue
        match = re.search(r'https?://[^\s<>")\]]+', field)
        if match:
            return match.group(0)

    fallback = (fallback_target or "http://TARGET").strip()
    if not fallback:
        return "http://TARGET"
    if re.match(r"^[a-zA-Z][a-zA-Z0-9+.-]*://", fallback):
        return fallback
    return f"http://{fallback}"


def _build_poc_content(
    finding: VulnerabilityFinding,
    target: str,
    filename: str,
) -> str:
    from jinja2 import Template

    normalized_type = _normalize_vuln_type(finding.vuln_type)
    poc_finding = finding.model_copy(update={"vuln_type": normalized_type})
    poc_body = PoCGenerator.generate_poc(poc_finding, target=target)

    template = Template(PYTHON_POC_TEMPLATE)
    return template.render(
        title=finding.title,
        severity=finding.severity,
        cve=finding.cve,
        vuln_type=normalized_type,
        target=target,
        generated_at=datetime.now().isoformat(),
        filename=filename,
        poc_body=poc_body.strip(),
    )


def generate_pocs(session: SessionState, output_dir: Path) -> list[Path]:
    """Generate PoC scripts for all findings in the session."""
    output_dir.mkdir(parents=True, exist_ok=True)
    generated: list[Path] = []

    for i, finding in enumerate(session.findings, 1):
        safe_name = _sanitize_filename_component(finding.title)
        cve_suffix = (
            f"_{_sanitize_filename_component(finding.cve, max_length=24)}" if finding.cve else ""
        )
        filename = f"poc_{i:02d}_{safe_name}{cve_suffix}.py"
        target = _extract_target(finding, session.target or "")
        content = _build_poc_content(finding, target=target, filename=filename)

        output_path = _next_available_poc_path(output_dir, filename)
        output_path.write_text(content, encoding="utf-8")
        generated.append(output_path)
        finding.poc_script = str(output_path)

    return generated


def generate_single_poc(
    title: str,
    severity: str = "High",
    cve: str | None = None,
    target: str = "http://TARGET",
    vuln_type: str = "",
) -> str:
    """Generate a single vulnerability-specific PoC script string."""
    finding = VulnerabilityFinding(
        title=title,
        severity=severity,
        cve=cve,
        vuln_type=vuln_type,
    )
    return _build_poc_content(
        finding,
        target=_extract_target(finding, target),
        filename=f"poc_{_normalize_vuln_type(vuln_type or 'custom')}.py",
    )


def _next_available_poc_path(output_dir: Path, filename: str) -> Path:
    """Return a writable PoC path, avoiding collisions with existing locked files."""
    base = output_dir / filename
    if not base.exists():
        return base

    stem = base.stem
    suffix = base.suffix
    for idx in range(1, 1000):
        candidate = output_dir / f"{stem}_{idx:02d}{suffix}"
        if not candidate.exists():
            return candidate
    return output_dir / f"{stem}_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}{suffix}"
