#!/usr/bin/env python3
from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import re
import shlex
import subprocess
import sys
import time
from collections import Counter
from pathlib import Path
from typing import Any


REPO_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_FOLD = REPO_ROOT / "build" / "release" / "src" / "cli" / "fold"
DEFAULT_CURL_SOURCE = REPO_ROOT.parent / "curl-source" / "curl"
DEFAULT_CURL_BINARY = Path("/usr/bin/curl")
RESULTS_ROOT = REPO_ROOT / "bench" / "results"

DEFAULT_SOURCE_DETECTORS = [
    "null-deref",
    "div-by-zero",
    "int-overflow",
    "oob",
    "int-truncation",
    "sign-conversion",
    "taint-track",
]
DEFAULT_BINARY_DETECTORS = ["security", "null-deref"]
SUCCESS_EXIT_CODES = {0, 1}  # fold returns 1 when findings are present.


def now_slug() -> str:
    return dt.datetime.now(dt.UTC).strftime("%Y-%m-%dT%H-%M-%SZ")


def split_csv(value: str) -> list[str]:
    return [item.strip() for item in value.split(",") if item.strip()]


def run_command(command: list[str], cwd: Path | None, timeout: int, stdout_path: Path,
                stderr_path: Path) -> dict[str, Any]:
    started = time.perf_counter()
    stdout_path.parent.mkdir(parents=True, exist_ok=True)
    timed_out = False
    try:
        with stdout_path.open("w") as stdout, stderr_path.open("w") as stderr:
            proc = subprocess.run(command, cwd=str(cwd) if cwd else None, stdout=stdout,
                                  stderr=stderr, text=True, timeout=timeout, check=False)
        exit_code = proc.returncode
    except subprocess.TimeoutExpired as exc:
        timed_out = True
        exit_code = 124
        stdout_path.write_text((exc.stdout or "") if isinstance(exc.stdout, str) else "")
        stderr_path.write_text((exc.stderr or "") if isinstance(exc.stderr, str) else "")
    elapsed = time.perf_counter() - started
    return {
        "command": command,
        "command_display": shlex.join(command),
        "cwd": str(cwd) if cwd else None,
        "exit_code": exit_code,
        "elapsed_s": elapsed,
        "timed_out": timed_out,
        "stdout": str(stdout_path),
        "stderr": str(stderr_path),
    }


def load_compile_db_files(compile_db: Path, max_files: int, path_regex: str | None) -> list[str]:
    db_file = compile_db / "compile_commands.json" if compile_db.is_dir() else compile_db
    if not db_file.exists():
        return []
    entries = json.loads(db_file.read_text())
    pattern = re.compile(path_regex) if path_regex else None
    files: list[str] = []
    seen: set[str] = set()
    for entry in entries:
        file_name = str(entry.get("file", ""))
        if not file_name or file_name in seen:
            continue
        if pattern and not pattern.search(file_name):
            continue
        seen.add(file_name)
        files.append(file_name)
    files.sort()
    if max_files > 0:
        files = files[:max_files]
    return files


def ensure_curl_compile_db(args: argparse.Namespace, out_dir: Path) -> tuple[Path | None, dict[str, Any] | None]:
    if args.compile_db:
        compile_db = Path(args.compile_db).expanduser().resolve()
        return compile_db, None

    curl_source = Path(args.curl_source).expanduser().resolve()
    build_dir = Path(args.source_build_dir).expanduser().resolve() if args.source_build_dir else out_dir / "curl-build"
    compile_db = build_dir / "compile_commands.json"
    if compile_db.exists():
        return compile_db, None

    if args.no_configure_source:
        return None, None

    command = [
        "cmake",
        "-S",
        str(curl_source),
        "-B",
        str(build_dir),
        "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON",
        "-DCMAKE_BUILD_TYPE=RelWithDebInfo",
    ]
    for cmake_arg in args.cmake_arg:
        command.append(cmake_arg)

    result = run_command(command, cwd=REPO_ROOT, timeout=args.configure_timeout,
                         stdout_path=out_dir / "curl-configure.stdout.log",
                         stderr_path=out_dir / "curl-configure.stderr.log")
    if compile_db.exists():
        return compile_db, result
    return None, result


def find_built_curl(build_dir: Path) -> Path | None:
    candidates = [
        build_dir / "src" / "curl",
        build_dir / "src" / "curl.exe",
        build_dir / "bin" / "curl",
    ]
    for candidate in candidates:
        if candidate.exists():
            return candidate.resolve()
    return None


def build_curl_binary(args: argparse.Namespace, out_dir: Path) -> tuple[Path | None, dict[str, Any] | None]:
    if not args.build_curl_binary:
        return None, None
    build_dir = Path(args.source_build_dir).expanduser().resolve() if args.source_build_dir else out_dir / "curl-build"
    command = ["cmake", "--build", str(build_dir), "--target", args.curl_target, "-j", str(args.curl_build_jobs)]
    result = run_command(command, cwd=REPO_ROOT, timeout=args.build_timeout,
                         stdout_path=out_dir / "curl-build.stdout.log",
                         stderr_path=out_dir / "curl-build.stderr.log")
    return find_built_curl(build_dir), result


def parse_sarif(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        return []
    doc = json.loads(path.read_text())
    findings: list[dict[str, Any]] = []
    for run in doc.get("runs", []):
        for result in run.get("results", []):
            props = result.get("properties", {})
            locations = result.get("locations", [])
            first_loc = locations[0] if locations else {}
            physical = first_loc.get("physicalLocation", {}) if isinstance(first_loc, dict) else {}
            artifact = physical.get("artifactLocation", {})
            region = physical.get("region", {})
            findings.append({
                "id": props.get("id", ""),
                "rule": str(result.get("ruleId", "fold.unknown")).removeprefix("fold."),
                "detector": props.get("detector", ""),
                "confidence": props.get("confidence", ""),
                "level": result.get("level", ""),
                "score": float(props.get("score", 0.0) or 0.0),
                "message": result.get("message", {}).get("text", ""),
                "file": artifact.get("uri", ""),
                "line": int(region.get("startLine", 0) or 0),
                "function": (first_loc.get("logicalLocations", [{}])[0] if first_loc.get("logicalLocations") else {}).get("fullyQualifiedName", ""),
                "binary": props.get("binary", {}) if isinstance(props.get("binary", {}), dict) else {},
                "suppressed": result.get("level") == "none" or bool(result.get("suppressions")),
            })
    return findings


def is_probably_text(path: Path) -> bool:
    return path.suffix.lower() in {
        ".c", ".cc", ".cpp", ".cxx", ".h", ".hh", ".hpp", ".hxx", ".in", ".pl",
        ".py", ".cmake", ".txt", ".md", ".rst",
    }


def source_candidates_for_string(root: Path, needle: str, limit: int = 8) -> list[dict[str, Any]]:
    if not needle:
        return []
    out: list[dict[str, Any]] = []
    skip_dirs = {".git", ".cache", "CMakeFiles", "Testing"}
    for base, dirs, files in os.walk(root):
        dirs[:] = [d for d in dirs if d not in skip_dirs]
        for name in files:
            path = Path(base) / name
            if not is_probably_text(path):
                continue
            try:
                if path.stat().st_size > 2_000_000:
                    continue
                for line_no, line in enumerate(path.read_text(errors="ignore").splitlines(), 1):
                    if needle in line:
                        out.append({"file": str(path), "line": line_no, "text": line.strip()})
                        if len(out) >= limit:
                            return out
            except OSError:
                continue
    return out


def source_candidates_for_symbol(root: Path, symbol: str, limit: int = 8) -> list[dict[str, Any]]:
    symbol = symbol.removeprefix("sym.").removeprefix("_").removeprefix("imp.")
    if not symbol:
        return []
    out: list[dict[str, Any]] = []
    pattern = re.compile(r"\b" + re.escape(symbol) + r"\b")
    skip_dirs = {".git", ".cache", "CMakeFiles", "Testing"}
    for base, dirs, files in os.walk(root):
        dirs[:] = [d for d in dirs if d not in skip_dirs]
        for name in files:
            path = Path(base) / name
            if not is_probably_text(path):
                continue
            try:
                if path.stat().st_size > 2_000_000:
                    continue
                for line_no, line in enumerate(path.read_text(errors="ignore").splitlines(), 1):
                    if pattern.search(line):
                        out.append({"file": str(path), "line": line_no, "text": line.strip()})
                        if len(out) >= limit:
                            return out
            except OSError:
                continue
    return out


def ranked_source_candidates(candidates: list[dict[str, Any]], source_paths: set[str],
                             limit: int = 8) -> list[dict[str, Any]]:
    seen: set[tuple[str, int]] = set()
    deduped: list[dict[str, Any]] = []
    for candidate in candidates:
        key = (candidate.get("file", ""), int(candidate.get("line", 0) or 0))
        if key in seen:
            continue
        seen.add(key)
        deduped.append(candidate)

    def rank(candidate: dict[str, Any]) -> tuple[int, int, int, str, int]:
        file = str(candidate.get("file", ""))
        normalized = file.replace(os.sep, "/")
        in_scan = 0 if file in source_paths else 1
        is_test = 1 if "/tests/" in normalized else 0
        is_source = 0 if "/src/" in normalized or "/lib/" in normalized else 1
        return (in_scan, is_test, is_source, file, int(candidate.get("line", 0) or 0))

    deduped.sort(key=rank)
    return deduped[:limit]


def summarize_findings(findings: list[dict[str, Any]]) -> dict[str, Any]:
    by_rule = Counter(finding["rule"] for finding in findings)
    by_detector = Counter(finding["detector"] or finding["rule"] for finding in findings)
    by_level = Counter(finding["level"] for finding in findings)
    scores = [finding["score"] for finding in findings]
    return {
        "total": len(findings),
        "suppressed": sum(1 for finding in findings if finding["suppressed"]),
        "by_rule": dict(sorted(by_rule.items())),
        "by_detector": dict(sorted(by_detector.items())),
        "by_level": dict(sorted(by_level.items())),
        "average_score": sum(scores) / len(scores) if scores else 0.0,
        "top_findings": sorted(findings, key=lambda item: item["score"], reverse=True)[:10],
    }


def run_source_scan(args: argparse.Namespace, out_dir: Path, compile_db: Path | None) -> dict[str, Any]:
    if args.skip_source:
        return {"skipped": True, "reason": "--skip-source"}
    if compile_db is None:
        return {"skipped": True, "reason": "no compile_commands.json available"}

    source_paths = load_compile_db_files(compile_db, args.source_max_files, args.source_path_regex)
    command = [
        str(Path(args.fold).expanduser().resolve()),
        "scan",
        "--compile-db",
        str(compile_db),
        "--format",
        "sarif",
        "--output",
        str(out_dir / "source.sarif"),
        "--processes",
        str(args.processes),
        "--budget-seconds",
        str(args.source_budget_seconds),
        "--per-query-ms",
        str(args.source_per_query_ms),
        "--per-tu-timeout",
        str(args.source_per_tu_timeout),
    ]
    for detector in split_csv(args.source_detectors):
        command.extend(["--detectors", detector])
    command.extend(source_paths)

    result = run_command(command, cwd=REPO_ROOT, timeout=args.source_timeout,
                         stdout_path=out_dir / "source.stdout.log",
                         stderr_path=out_dir / "source.stderr.log")
    findings = parse_sarif(out_dir / "source.sarif")
    result.update({
        "skipped": False,
        "compile_db": str(compile_db),
        "source_paths_count": len(source_paths),
        "source_paths": source_paths,
        "report": str(out_dir / "source.sarif"),
        "success": result["exit_code"] in SUCCESS_EXIT_CODES and (out_dir / "source.sarif").exists(),
        "findings": summarize_findings(findings),
    })
    return result


def run_binary_scan(args: argparse.Namespace, out_dir: Path, binary_path: Path) -> dict[str, Any]:
    if args.skip_binary:
        return {"skipped": True, "reason": "--skip-binary"}
    if not binary_path.exists():
        return {"skipped": True, "reason": f"binary does not exist: {binary_path}"}

    command = [
        str(Path(args.fold).expanduser().resolve()),
        "scan",
        "--input-type",
        "binary",
        "--binary",
        str(binary_path),
        "--binary-max-functions",
        str(args.binary_max_functions),
        "--format",
        "sarif",
        "--output",
        str(out_dir / "binary.sarif"),
    ]
    for detector in split_csv(args.binary_detectors):
        command.extend(["--detectors", detector])

    result = run_command(command, cwd=REPO_ROOT, timeout=args.binary_timeout,
                         stdout_path=out_dir / "binary.stdout.log",
                         stderr_path=out_dir / "binary.stderr.log")
    findings = parse_sarif(out_dir / "binary.sarif")
    result.update({
        "skipped": False,
        "binary": str(binary_path),
        "report": str(out_dir / "binary.sarif"),
        "success": result["exit_code"] in SUCCESS_EXIT_CODES and (out_dir / "binary.sarif").exists(),
        "findings": summarize_findings(findings),
    })
    return result


def comparison(source: dict[str, Any], binary: dict[str, Any]) -> dict[str, Any]:
    if source.get("skipped") or binary.get("skipped"):
        return {"available": False}
    source_rules = set(source.get("findings", {}).get("by_rule", {}).keys())
    binary_rules = set(binary.get("findings", {}).get("by_rule", {}).keys())
    source_detectors = set(source.get("findings", {}).get("by_detector", {}).keys())
    binary_detectors = set(binary.get("findings", {}).get("by_detector", {}).keys())
    return {
        "available": True,
        "source_total": source["findings"]["total"],
        "binary_total": binary["findings"]["total"],
        "common_rules": sorted(source_rules & binary_rules),
        "source_only_rules": sorted(source_rules - binary_rules),
        "binary_only_rules": sorted(binary_rules - source_rules),
        "common_detectors": sorted(source_detectors & binary_detectors),
        "source_only_detectors": sorted(source_detectors - binary_detectors),
        "binary_only_detectors": sorted(binary_detectors - source_detectors),
    }


def correlate_binary_source(summary: dict[str, Any], source_findings: list[dict[str, Any]],
                            binary_findings: list[dict[str, Any]]) -> list[dict[str, Any]]:
    root = Path(summary["inputs"]["curl_source"])
    source_paths = set(summary.get("source", {}).get("source_paths", []))
    by_file: dict[str, list[dict[str, Any]]] = {}
    for finding in source_findings:
        by_file.setdefault(finding.get("file", ""), []).append(finding)

    correlated: list[dict[str, Any]] = []
    for finding in binary_findings:
        binary = finding.get("binary", {})
        matched = str(binary.get("matched_string", "") or "")
        xref_function = str(binary.get("xref_function", "") or finding.get("function", "") or "")
        candidates = ranked_source_candidates(
            source_candidates_for_symbol(root, xref_function, limit=16) +
            source_candidates_for_string(root, matched, limit=32),
            source_paths,
        )
        candidate_files = {candidate["file"] for candidate in candidates}
        related_source = []
        for file in sorted(candidate_files):
            for source_finding in by_file.get(file, [])[:5]:
                related_source.append(source_finding)
        correlated.append({
            "binary_finding": finding,
            "matched_string": matched,
            "xref_function": xref_function,
            "source_candidates": candidates,
            "candidate_files_in_source_scan": sorted(candidate_files & source_paths),
            "candidate_files_outside_source_scan": sorted(candidate_files - source_paths),
            "source_findings_in_candidate_files": related_source,
        })
    return correlated


def markdown_report(summary: dict[str, Any]) -> str:
    lines = ["# curl source-vs-binary A/B", ""]
    lines.append(f"- timestamp: `{summary['timestamp']}`")
    lines.append(f"- curl source: `{summary['inputs']['curl_source']}`")
    lines.append(f"- curl binary: `{summary['inputs']['curl_binary']}`")
    lines.append("")

    for name in ("source", "binary"):
        run = summary[name]
        lines.append(f"## {name.title()} Scan")
        if run.get("skipped"):
            lines.append(f"Skipped: {run.get('reason', 'unknown')}")
            lines.append("")
            continue
        lines.append(f"- success: `{run.get('success')}`")
        lines.append(f"- exit code: `{run.get('exit_code')}`")
        lines.append(f"- elapsed: `{run.get('elapsed_s', 0.0):.2f}s`")
        lines.append(f"- command: `{run.get('command_display', '')}`")
        findings = run.get("findings", {})
        lines.append(f"- findings: `{findings.get('total', 0)}`")
        lines.append(f"- suppressed: `{findings.get('suppressed', 0)}`")
        lines.append(f"- by detector: `{json.dumps(findings.get('by_detector', {}), sort_keys=True)}`")
        lines.append("")

    comp = summary.get("comparison", {})
    lines.append("## Comparison")
    if not comp.get("available"):
        lines.append("Comparison unavailable because one side was skipped.")
    else:
        lines.append(f"- source findings: `{comp['source_total']}`")
        lines.append(f"- binary findings: `{comp['binary_total']}`")
        lines.append(f"- common rules: `{', '.join(comp['common_rules']) or '<none>'}`")
        lines.append(f"- source-only rules: `{', '.join(comp['source_only_rules']) or '<none>'}`")
        lines.append(f"- binary-only rules: `{', '.join(comp['binary_only_rules']) or '<none>'}`")
    lines.append("")

    correlations = summary.get("correlation", [])
    if correlations:
        lines.append("## Binary/Source Correlation")
        lines.append("| Binary finding | Binary evidence | Source candidates | Source findings |")
        lines.append("|---|---|---|---|")
        for item in correlations:
            finding = item.get("binary_finding", {})
            binary = finding.get("binary", {})
            offset = binary.get("offset")
            vaddr = binary.get("vaddr")
            section = binary.get("section") or ""
            xref = item.get("xref_function") or ""
            evidence = []
            if offset is not None:
                evidence.append(f"offset `0x{int(offset):x}`")
            if vaddr is not None:
                evidence.append(f"vaddr `0x{int(vaddr):x}`")
            if section:
                evidence.append(f"section `{section}`")
            if xref:
                evidence.append(f"xref `{xref}`")
            candidates = []
            for candidate in item.get("source_candidates", [])[:4]:
                candidates.append(f"`{candidate['file']}:{candidate['line']}`")
            source_findings = []
            for source_finding in item.get("source_findings_in_candidate_files", [])[:4]:
                source_findings.append(
                    f"`{source_finding.get('rule', '')}` at `{source_finding.get('file', '')}:{source_finding.get('line', 0)}`"
                )
            if not source_findings and item.get("candidate_files_outside_source_scan"):
                source_findings.append("candidate file was outside this source scan")
            lines.append(
                f"| `{finding.get('detector', finding.get('rule', ''))}` | "
                f"{'; '.join(evidence) or '<none>'} | "
                f"{'<br>'.join(candidates) or '<none>'} | "
                f"{'<br>'.join(source_findings) or '<none>'} |"
            )
        lines.append("")
    return "\n".join(lines)


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run fold source-vs-binary A/B on curl.")
    parser.add_argument("--fold", default=str(DEFAULT_FOLD))
    parser.add_argument("--curl-source", default=str(DEFAULT_CURL_SOURCE))
    parser.add_argument("--curl-binary", default=str(DEFAULT_CURL_BINARY))
    parser.add_argument("--compile-db", default="")
    parser.add_argument("--source-build-dir", default="")
    parser.add_argument("--out-dir", default="")
    parser.add_argument("--source-detectors", default=",".join(DEFAULT_SOURCE_DETECTORS))
    parser.add_argument("--binary-detectors", default=",".join(DEFAULT_BINARY_DETECTORS))
    parser.add_argument("--source-max-files", type=int, default=0)
    parser.add_argument("--source-path-regex", default="")
    parser.add_argument("--binary-max-functions", type=int, default=64)
    parser.add_argument("--processes", type=int, default=1)
    parser.add_argument("--source-budget-seconds", type=int, default=30)
    parser.add_argument("--source-per-query-ms", type=int, default=500)
    parser.add_argument("--source-per-tu-timeout", type=int, default=30)
    parser.add_argument("--skip-source", action="store_true")
    parser.add_argument("--skip-binary", action="store_true")
    parser.add_argument("--no-configure-source", action="store_true")
    parser.add_argument("--build-curl-binary", action="store_true")
    parser.add_argument("--curl-target", default="curl")
    parser.add_argument("--curl-build-jobs", type=int, default=os.cpu_count() or 2)
    parser.add_argument("--cmake-arg", action="append", default=[])
    parser.add_argument("--configure-timeout", type=int, default=180)
    parser.add_argument("--build-timeout", type=int, default=600)
    parser.add_argument("--source-timeout", type=int, default=600)
    parser.add_argument("--binary-timeout", type=int, default=300)
    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    out_dir = Path(args.out_dir).expanduser().resolve() if args.out_dir else (
        RESULTS_ROOT / f"{now_slug()}_curl-ab"
    )
    out_dir.mkdir(parents=True, exist_ok=True)

    compile_db, configure_result = ensure_curl_compile_db(args, out_dir)
    built_binary, build_result = build_curl_binary(args, out_dir)
    binary_path = built_binary or Path(args.curl_binary).expanduser().resolve()

    source = run_source_scan(args, out_dir, compile_db)
    binary = run_binary_scan(args, out_dir, binary_path)
    source_findings = parse_sarif(out_dir / "source.sarif")
    binary_findings = parse_sarif(out_dir / "binary.sarif")

    summary = {
        "timestamp": now_slug(),
        "inputs": {
            "fold": str(Path(args.fold).expanduser().resolve()),
            "curl_source": str(Path(args.curl_source).expanduser().resolve()),
            "curl_binary": str(binary_path),
            "compile_db": str(compile_db) if compile_db else None,
        },
        "configure": configure_result,
        "curl_build": build_result,
        "source": source,
        "binary": binary,
        "comparison": comparison(source, binary),
    }
    summary["correlation"] = correlate_binary_source(summary, source_findings, binary_findings)

    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    (out_dir / "summary.md").write_text(markdown_report(summary))
    print(f"wrote {out_dir / 'summary.md'}")

    source_ok = source.get("skipped") or source.get("success")
    binary_ok = binary.get("skipped") or binary.get("success")
    return 0 if source_ok and binary_ok else 1


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))