#!/usr/bin/env python3
"""Add an org id to Clinic API feature-gate allowlists."""

from __future__ import annotations

import argparse
from pathlib import Path
import re
import sys


DEFAULT_GATES = ("orgs", "slow-query", "top-sql")
CLINIC_DIR = Path(__file__).resolve().parents[1]
DEFAULT_CONFIG = CLINIC_DIR / "clinic-api/envs/prod-cloud/configmap/config.yaml"
WETECH_ORG_ID = "1372813089209061633"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Add an org id to Clinic API external-feature-gates enabledOrgs lists."
    )
    parser.add_argument("--org-id", required=True, help="TiDB Cloud org/tenant id.")
    parser.add_argument(
        "--comment",
        required=True,
        help="Inline comment to write after the org id, for example: Wetech.",
    )
    parser.add_argument(
        "--config",
        type=Path,
        default=DEFAULT_CONFIG,
        help=f"Config file to edit. Default: {DEFAULT_CONFIG}",
    )
    parser.add_argument(
        "--gates",
        nargs="+",
        default=list(DEFAULT_GATES),
        help=f"Feature gate names to update. Default: {' '.join(DEFAULT_GATES)}",
    )
    parser.add_argument(
        "--after-org-id",
        default=WETECH_ORG_ID,
        help="Prefer inserting after this org id when it exists in the target list.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print what would change without writing the file.",
    )
    return parser.parse_args()


def find_gate_block(lines: list[str], gate: str) -> tuple[int, int]:
    gate_pattern = re.compile(rf"^  {re.escape(gate)}:\s*$")
    start = next((idx for idx, line in enumerate(lines) if gate_pattern.match(line)), -1)
    if start < 0:
        raise ValueError(f"feature gate {gate!r} was not found")

    end = len(lines)
    for idx in range(start + 1, len(lines)):
        if re.match(r"^  [A-Za-z0-9_-]+:\s*$", lines[idx]):
            end = idx
            break
    return start, end


def find_enabled_orgs_list(lines: list[str], gate: str) -> tuple[int, int]:
    block_start, block_end = find_gate_block(lines, gate)
    enabled_idx = next(
        (
            idx
            for idx in range(block_start + 1, block_end)
            if re.match(r"^    enabledOrgs:\s*$", lines[idx])
        ),
        -1,
    )
    if enabled_idx < 0:
        raise ValueError(f"feature gate {gate!r} does not contain enabledOrgs")

    list_start = enabled_idx + 1
    list_end = list_start
    while list_end < block_end and re.match(r"^      - ", lines[list_end]):
        list_end += 1
    return list_start, list_end


def add_org_to_gate(
    lines: list[str], gate: str, org_id: str, comment: str, after_org_id: str
) -> bool:
    list_start, list_end = find_enabled_orgs_list(lines, gate)
    org_pattern = re.compile(rf"^\s*-\s*['\"]?{re.escape(org_id)}['\"]?\b")

    for idx in range(list_start, list_end):
        if org_pattern.search(lines[idx]):
            return False

    insert_at = list_end
    after_pattern = re.compile(rf"^\s*-\s*['\"]?{re.escape(after_org_id)}['\"]?\b")
    for idx in range(list_start, list_end):
        if after_pattern.search(lines[idx]):
            insert_at = idx + 1
            break

    lines.insert(insert_at, f"      - '{org_id}' # {comment}\n")
    return True


def main() -> int:
    args = parse_args()
    if not re.fullmatch(r"\d+", args.org_id):
        print(f"invalid --org-id {args.org_id!r}: expected digits only", file=sys.stderr)
        return 2

    config = args.config
    if not config.exists():
        print(f"config file not found: {config}", file=sys.stderr)
        return 2

    lines = config.read_text().splitlines(keepends=True)
    changed_gates: list[str] = []
    for gate in args.gates:
        if add_org_to_gate(lines, gate, args.org_id, args.comment, args.after_org_id):
            changed_gates.append(gate)

    if args.dry_run:
        if changed_gates:
            print(f"would add {args.org_id} to: {', '.join(changed_gates)}")
        else:
            print(f"{args.org_id} is already present in: {', '.join(args.gates)}")
        return 0

    if changed_gates:
        config.write_text("".join(lines))
        print(f"added {args.org_id} to: {', '.join(changed_gates)}")
    else:
        print(f"{args.org_id} is already present in: {', '.join(args.gates)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())