#!/usr/bin/env python3
import argparse, json, re, subprocess, sys
from pathlib import Path


def get_checkov_comments(repo, pr_number, comment_index=None):
    result = subprocess.run(
        ["gh", "pr", "view", str(pr_number), "--repo", repo, "--json", "comments"],
        capture_output=True, text=True, check=True
    )
    comments = json.loads(result.stdout).get("comments", [])
    checkov_comments = [
        c for c in comments
        if ("checkov-scan" in c.get("author", {}).get("login", "") or "Checkov notification" in c.get("body", ""))
        and "Checkov check FAILED:" in c.get("body", "")
    ]
    if not checkov_comments:
        return []
    if comment_index is not None:
        try:
            return [checkov_comments[comment_index]["body"]]
        except IndexError:
            print(f"Comment index {comment_index} out of range ({len(checkov_comments)} checkov comment(s) found)")
            return []
    return [c["body"] for c in checkov_comments]


def parse_checkov_failures(body, debug=False):
    failures = []
    # Handle both unchecked [ ] and checked [x] checkboxes (e.g. on merged PRs)
    blocks = re.split(r'(?=- \[[ x]\] `?Checkov check FAILED:)', body)
    if debug:
        print(f"DEBUG: split into {len(blocks)} block(s)")
    for i, block in enumerate(blocks):
        check_match = re.search(r'Checkov check FAILED:\s+([A-Z][A-Z0-9_]+)', block)
        # Description: #### heading with optional bold markers
        desc_match  = re.search(r'(?:####\s+)?\*\*([^*\n]{5,})\*\*', block)
        # Resource/File headers may or may not use **bold**
        res_match   = re.search(r'\|\s+\*{0,2}Resource\*{0,2}\s+\|\s+(\S+)\s+\|', block)
        # File path may contain " -> " for module references — take the first path only
        file_match  = re.search(r'\|\s+\*{0,2}File\*{0,2}\s+\|\s+([^\s|]+?)(?:\s*->.*?)?:\d[\d\-]*\s+\|', block)
        if debug:
            print(f"DEBUG block {i}: check={check_match.group(1) if check_match else None} "
                  f"res={res_match.group(1) if res_match else None} "
                  f"file={file_match.group(1) if file_match else None}")
        if not (check_match and res_match and file_match):
            continue
        failures.append({
            "check_id":    check_match.group(1),
            "description": desc_match.group(1).strip() if desc_match else "",
            "resource":    res_match.group(1).strip(),
            "file_path":   file_match.group(1).strip(),
        })
    return failures


def remap_path(file_path, path_maps):
    for remote_prefix, local_prefix in path_maps:
        if file_path.startswith(remote_prefix):
            return (local_prefix + file_path[len(remote_prefix):]).lstrip("/")
    return file_path


def find_resource_line(lines, resource_address):
    parts = resource_address.split(".")
    if len(parts) < 2:
        return None

    # module.module_name.resource_type.resource_name[...]
    if parts[0] == "module":
        module_name = parts[1]
        pattern = re.compile(rf'^\s*module\s+"{re.escape(module_name)}"\s*{{')
        for i, line in enumerate(lines):
            if pattern.match(line):
                return i
        return None

    res_type, res_name = parts[0], parts[1]
    for block_keyword in ("data", "resource"):
        pattern = re.compile(
            rf'^\s*{block_keyword}\s+"{re.escape(res_type)}"\s+"{re.escape(res_name)}"\s*{{'
        )
        for i, line in enumerate(lines):
            if pattern.match(line):
                return i
    return None


def find_module_file(resource, root):
    """For module.x.y resources, search all local .tf files for the module block."""
    parts = resource.split(".")
    if parts[0] != "module" or len(parts) < 2:
        return None, None
    module_name = parts[1]
    pattern = re.compile(rf'^\s*module\s+"{re.escape(module_name)}"\s*{{')
    for tf_file in sorted(Path(root).rglob("*.tf")):
        # skip cached module dirs
        if ".terraform" in tf_file.parts:
            continue
        lines = tf_file.read_text().splitlines(keepends=True)
        for i, line in enumerate(lines):
            if pattern.match(line):
                return tf_file, lines
    return None, None


def add_skip(file_path, resource, check_id, description, root, path_maps):
    mapped = remap_path(file_path, path_maps)
    full_path = Path(root) / mapped

    # If the reported file is inside a cached .terraform-modules dir,
    # find the local module call instead.
    if not full_path.exists() or ".terraform-modules" in file_path:
        parts = resource.split(".")
        if parts[0] == "module":
            local_file, lines = find_module_file(resource, root)
            if local_file is None:
                print(f"  MODULE NOT FOUND locally: {resource}")
                return False
            full_path = local_file
            mapped = str(local_file.relative_to(Path(root)))
        else:
            print(f"  NOT FOUND: {full_path}")
            return False
    else:
        lines = full_path.read_text().splitlines(keepends=True)

    idx = find_resource_line(lines, resource)
    if idx is None:
        print(f"  RESOURCE NOT FOUND: {resource} in {mapped}")
        return False

    # Scan until the closing } of the resource/data block
    depth = 0
    block_lines = []
    for line in lines[idx:]:
        depth += line.count('{') - line.count('}')
        block_lines.append(line)
        if depth <= 0:
            break
    if f'checkov:skip={check_id}' in ''.join(block_lines):
        print(f"  ALREADY PRESENT: {check_id} on {resource}")
        return False

    indent = re.match(r'^(\s*)', lines[idx]).group(1)
    lines.insert(idx + 1, f"{indent}  #checkov:skip={check_id}: {description}\n")
    full_path.write_text(''.join(lines))
    print(f"  ADDED: {check_id} -> {resource} ({mapped}:{idx + 1})")
    return True


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("pr", type=int)
    parser.add_argument("--repo", required=True)
    parser.add_argument("--root", default=".")
    parser.add_argument(
        "--path-map", action="append", default=[], metavar="REMOTE=LOCAL",
        help="Map remote path prefix to local. Can be repeated. e.g. aoma-core/prod/iam/core=sme-aoma-core-prod/iam"
    )
    parser.add_argument("--debug", action="store_true", help="Print raw parsing debug info")
    parser.add_argument("--comment-index", type=int, default=None,
                        help="Index of a specific checkov comment (0=oldest, -1=latest). Default: use all comments.")
    parser.add_argument("--list-comments", action="store_true",
                        help="List all checkov comments with their index and exit")
    args = parser.parse_args()

    path_maps = []
    for m in args.path_map:
        if "=" not in m:
            print(f"Invalid --path-map '{m}', expected REMOTE=LOCAL"); sys.exit(1)
        remote, local = m.split("=", 1)
        path_maps.append((remote, local))

    print(f"Fetching PR #{args.pr} comments from {args.repo}...")

    if args.list_comments:
        result = subprocess.run(
            ["gh", "pr", "view", str(args.pr), "--repo", args.repo, "--json", "comments"],
            capture_output=True, text=True, check=True
        )
        all_comments = json.loads(result.stdout).get("comments", [])
        checkov = [c for c in all_comments
                   if ("checkov-scan" in c.get("author", {}).get("login", "") or "Checkov notification" in c.get("body", ""))
                   and "Checkov check FAILED:" in c.get("body", "")]
        print(f"Found {len(checkov)} checkov comment(s):")
        for i, c in enumerate(checkov):
            preview = c["body"][:80].replace("\n", " ")
            print(f"  [{i}] (also [{i - len(checkov)}]) {c.get('createdAt', '')} — {preview}")
        sys.exit(0)

    bodies = get_checkov_comments(args.repo, args.pr, comment_index=args.comment_index)
    if not bodies:
        print("No checkov-scan comment found."); sys.exit(0)

    seen = set()
    failures = []
    for body in bodies:
        for f in parse_checkov_failures(body, debug=args.debug):
            key = (f["check_id"], f["resource"], f["file_path"])
            if key not in seen:
                seen.add(key)
                failures.append(f)

    if not failures:
        print("No failures parsed from comment."); sys.exit(0)

    print(f"Found {len(failures)} unique failure(s) across {len(bodies)} comment(s).\n")
    added = 0
    for f in failures:
        print(f"[{f['check_id']}] {f['resource']} ({f['file_path']})")
        if add_skip(f["file_path"], f["resource"], f["check_id"], f["description"], args.root, path_maps):
            added += 1

    print(f"\nDone. {added} annotation(s) added.")


if __name__ == "__main__":
    main()