#!/usr/bin/env python3
"""Locate alert rules under a local runbooks/rules checkout by alert name."""

import argparse
import json
import os
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

ALERT_RE = re.compile(
    r'(?m)^(?P<indent>\s*)(?:-\s*)?alert\s*:\s*(?P<quote>["\']?)(?P<name>[^"\n#]+?)(?P=quote)\s*$'
)
EXPR_RE = re.compile(r"(?m)^(?P<indent>\s*)expr\s*:\s*(?P<expr>.+?)\s*$")
FOR_RE = re.compile(r"(?m)^\s*for\s*:\s*(?P<duration>.+?)\s*$")
TOKEN_RE = re.compile(r"[a-z0-9]+")


@dataclass
class Match:
    file: str
    alert: str
    match_type: str
    score: float
    context_score: float
    expr: str | None
    duration: str | None
    snippet: str

    def to_dict(self) -> dict[str, object]:
        return {
            "file": self.file,
            "alert": self.alert,
            "match_type": self.match_type,
            "score": round(self.score, 4),
            "context_score": round(self.context_score, 4),
            "expr": self.expr,
            "for": self.duration,
            "snippet": self.snippet,
        }


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--alert-name", required=True, help="Alert name to search for")
    parser.add_argument("--component", help="Optional alert component for disambiguation")
    parser.add_argument(
        "--message",
        action="append",
        default=[],
        help="Optional alert message text for disambiguation; may be repeated",
    )
    parser.add_argument(
        "--runbooks-dir",
        help="Local path to the runbooks repo root; must contain rules/",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=10,
        help="Maximum fuzzy candidates to return (default: 10)",
    )
    return parser.parse_args(argv)


def normalize(value: str) -> str:
    return "".join(ch for ch in value.lower() if ch.isalnum())


def tokenize(value: str) -> set[str]:
    return set(TOKEN_RE.findall(value.lower()))


def candidate_runbooks_dirs(explicit: str | None) -> Iterable[Path]:
    seen: set[Path] = set()
    raw_candidates = [
        explicit,
        os.environ.get("TIDBCLOUD_RUNBOOKS_DIR"),
        "runbooks",
        "../runbooks",
        "../../runbooks",
    ]
    for raw in raw_candidates:
        if not raw:
            continue
        path = Path(raw).expanduser().resolve()
        if path in seen:
            continue
        seen.add(path)
        yield path


def resolve_runbooks_dir(explicit: str | None) -> Path:
    checked: list[str] = []
    for candidate in candidate_runbooks_dirs(explicit):
        checked.append(str(candidate))
        if (candidate / "rules").is_dir():
            return candidate
    raise FileNotFoundError(
        "Could not find a local runbooks checkout with a rules/ directory. "
        f"Checked: {', '.join(checked) if checked else '(none)'}"
    )


def extract_expr(block: str) -> str | None:
    expr_match = EXPR_RE.search(block)
    if not expr_match:
        return None

    inline_expr = expr_match.group("expr").strip()
    if inline_expr not in {"|", ">"}:
        return inline_expr

    lines = block.splitlines()
    for idx, line in enumerate(lines):
        if line == expr_match.group(0):
            expr_indent = len(expr_match.group("indent"))
            collected: list[str] = []
            for follow in lines[idx + 1 :]:
                stripped = follow.strip()
                current_indent = len(follow) - len(follow.lstrip(" "))
                if stripped and current_indent <= expr_indent:
                    break
                if not stripped:
                    collected.append("")
                    continue
                collected.append(follow[expr_indent + 2 :].rstrip())
            expr = "\n".join(collected).strip()
            return expr or inline_expr
    return inline_expr


def extract_rule_block(text: str, alert_match: re.Match[str]) -> tuple[str | None, str | None, str]:
    start = alert_match.start()
    next_match = ALERT_RE.search(text, alert_match.end())
    end = next_match.start() if next_match else len(text)
    block = text[start:end].strip("\n")
    for_match = FOR_RE.search(block)
    lines = [line.rstrip() for line in block.splitlines()]
    snippet = "\n".join(lines[:12]).strip()
    return (
        extract_expr(block),
        for_match.group("duration").strip() if for_match else None,
        snippet,
    )


def score_match(target: str, candidate: str) -> tuple[str | None, float]:
    if candidate == target:
        return "exact", 1.0
    if candidate.lower() == target.lower():
        return "case_insensitive_exact", 0.99
    if normalize(candidate) == normalize(target):
        return "normalized_exact", 0.98
    candidate_lower = candidate.lower()
    target_lower = target.lower()
    if target_lower in candidate_lower or candidate_lower in target_lower:
        shared = min(len(candidate_lower), len(target_lower))
        longest = max(len(candidate_lower), len(target_lower), 1)
        return "substring", 0.75 + (shared / longest) * 0.15

    candidate_tokens = tokenize(candidate)
    target_tokens = tokenize(target)
    if candidate_tokens and target_tokens:
        overlap = candidate_tokens & target_tokens
        if overlap:
            union = candidate_tokens | target_tokens
            return "token_overlap", 0.45 + (len(overlap) / len(union)) * 0.25
    return None, 0.0


def score_context(file_path: str, rule_block: str, component: str | None, messages: list[str]) -> float:
    score = 0.0
    haystacks = [file_path.lower(), rule_block.lower()]
    if component:
        component_lower = component.strip().lower()
        if component_lower:
            if any(component_lower in haystack for haystack in haystacks):
                score += 0.4
            else:
                component_tokens = tokenize(component_lower)
                block_tokens = tokenize(" ".join(haystacks))
                if component_tokens and (component_tokens & block_tokens):
                    score += 0.2

    block_tokens = tokenize(" ".join(haystacks))
    for message in messages:
        msg = (message or "").strip().lower()
        if not msg:
            continue
        if msg in haystacks[1]:
            score += 0.35
            continue
        msg_tokens = tokenize(msg)
        if not msg_tokens:
            continue
        overlap = msg_tokens & block_tokens
        if overlap:
            score += min(0.3, 0.06 * len(overlap))
    return score


def search_rules(
    runbooks_dir: Path,
    alert_name: str,
    component: str | None,
    messages: list[str],
    limit: int,
) -> dict[str, object]:
    rules_dir = runbooks_dir / "rules"
    exact_matches: list[Match] = []
    fuzzy_matches: list[Match] = []
    searched_files = 0

    for file_path in sorted(rules_dir.rglob("*")):
        if not file_path.is_file() or file_path.suffix.lower() not in {".yml", ".yaml"}:
            continue
        searched_files += 1
        try:
            text = file_path.read_text(encoding="utf-8")
        except UnicodeDecodeError:
            text = file_path.read_text(encoding="utf-8", errors="replace")

        for alert_match in ALERT_RE.finditer(text):
            candidate_alert = alert_match.group("name").strip()
            match_type, score = score_match(alert_name, candidate_alert)
            if not match_type:
                continue

            expr, duration, snippet = extract_rule_block(text, alert_match)
            context_score = score_context(
                str(file_path.relative_to(runbooks_dir)),
                snippet,
                component,
                messages,
            )
            match = Match(
                file=str(file_path.relative_to(runbooks_dir)),
                alert=candidate_alert,
                match_type=match_type,
                score=score,
                context_score=context_score,
                expr=expr,
                duration=duration,
                snippet=snippet,
            )
            if match_type in {"exact", "case_insensitive_exact", "normalized_exact"}:
                exact_matches.append(match)
            else:
                fuzzy_matches.append(match)

    fuzzy_matches.sort(key=lambda item: (-(item.score + item.context_score), -item.score, item.file, item.alert))
    exact_matches.sort(key=lambda item: (-(item.score + item.context_score), -item.context_score, item.file, item.alert))

    ranked = exact_matches or fuzzy_matches
    top_match = ranked[0].to_dict() if ranked else None

    return {
        "alert_name": alert_name,
        "component": component,
        "messages": messages,
        "runbooks_dir": str(runbooks_dir),
        "searched_files": searched_files,
        "top_match": top_match,
        "exact_matches": [match.to_dict() for match in exact_matches],
        "fuzzy_matches": [match.to_dict() for match in fuzzy_matches[: max(limit, 0)]],
    }


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    try:
        runbooks_dir = resolve_runbooks_dir(args.runbooks_dir)
        result = search_rules(
            runbooks_dir,
            args.alert_name,
            args.component,
            args.message,
            args.limit,
        )
    except FileNotFoundError as exc:
        print(json.dumps({"status": "error", "error": str(exc)}))
        return 2
    except Exception as exc:  # pragma: no cover
        print(json.dumps({"status": "error", "error": repr(exc)}))
        return 1

    print(json.dumps({"status": "ok", **result}, ensure_ascii=False, indent=2))
    return 0


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