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

import argparse
import glob
import json
import os
import re
import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from datetime import date, datetime, timedelta
from pathlib import Path


@dataclass
class CleanupAction:
    category: str
    path: str
    size_bytes: int
    reason: str
    operation: str


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Cleanup host-level runtime logs and report local runtime disk pressure.")
    parser.add_argument("--dry-run", action="store_true", help="report only; do not mutate files")
    parser.add_argument("--summary-json", help="optional summary json output path")
    parser.add_argument(
        "--max-log-bytes",
        type=int,
        default=512 * 1024 * 1024,
        help="truncate managed runtime log files above this size",
    )
    parser.add_argument(
        "--journal-vacuum-size",
        default="200M",
        help="systemd journal target size passed to journalctl --vacuum-size",
    )
    parser.add_argument(
        "--shared-tidb-container-name",
        default="nexuscrm-shared-tidb",
        help="local shared TiDB docker container to report",
    )
    parser.add_argument(
        "--shared-tidb-image",
        default="pingcap/tidb:v8.5.0",
        help="image used when recreating the shared TiDB docker container",
    )
    parser.add_argument(
        "--recreate-shared-tidb",
        action="store_true",
        help="remove and recreate the local shared TiDB docker container",
    )
    parser.add_argument(
        "--runner-temp-max-age-hours",
        type=int,
        default=24,
        help="remove idle runner _work temp/action/checkouts older than this many hours",
    )
    parser.add_argument(
        "--runner-diag-max-age-days",
        type=int,
        default=7,
        help="remove runner _diag files older than this many days",
    )
    parser.add_argument(
        "--managed-temp-max-age-hours",
        type=int,
        default=6,
        help="remove managed host temp directories older than this many hours",
    )
    parser.add_argument(
        "--sf-export-keep-days",
        type=int,
        default=5,
        help="keep this many date-named Salesforce export directories",
    )
    parser.add_argument(
        "--log-keep-runs",
        type=int,
        default=10,
        help="keep this many newest run-scoped backend log directories per staging/test slot",
    )
    return parser.parse_args()


def log(message: str) -> None:
    print(message, file=sys.stderr)


def format_bytes(size_bytes: int) -> str:
    units = ["B", "KB", "MB", "GB", "TB"]
    size = float(size_bytes)
    for unit in units:
        if size < 1024 or unit == units[-1]:
            if unit == "B":
                return f"{int(size)}{unit}"
            return f"{size:.1f}{unit}"
        size /= 1024
    return f"{size_bytes}B"


def file_size(path: str) -> int:
    try:
        return os.stat(path).st_size
    except OSError:
        return 0


def path_size(path: Path) -> int:
    if not path.exists():
        return 0
    if path.is_file() or path.is_symlink():
        return file_size(str(path))
    total = 0
    for root, dirs, files in os.walk(path):
        root_path = Path(root)
        for name in dirs:
            child = root_path / name
            if child.is_symlink():
                total += file_size(str(child))
        for name in files:
            total += file_size(str(root_path / name))
    return total


def older_than(path: Path, cutoff_epoch: float) -> bool:
    try:
        return path.stat().st_mtime < cutoff_epoch
    except OSError:
        return False


def managed_log_paths() -> list[str]:
    patterns = [
        "/var/lib/docker/containers/*/*-json.log",
        "/home/ec2-user/nexuscrm-deployment/offline/envs/*/logs/backend.*.log",
        "/home/ec2-user/nexuscrm-deployment/offline/envs/staging/*/logs/backend.*.log",
        "/home/ec2-user/nexuscrm-deployment/offline/envs/test/*/logs/backend.*.log",
        "/home/ec2-user/nexuscrm-deployment/offline/envs/*/logs/run-*/backend.*.log",
        "/home/ec2-user/nexuscrm-deployment/offline/envs/staging/*/logs/run-*/backend.*.log",
        "/home/ec2-user/nexuscrm-deployment/offline/envs/test/*/logs/run-*/backend.*.log",
    ]
    paths: set[str] = set()
    for pattern in patterns:
        paths.update(glob.glob(pattern))
    return sorted(paths)


def run_scoped_log_dir_roots() -> list[str]:
    roots = [
        "/home/ec2-user/nexuscrm-deployment/offline/envs/staging/logs",
        "/home/ec2-user/nexuscrm-deployment/offline/envs/staging/blue/logs",
        "/home/ec2-user/nexuscrm-deployment/offline/envs/staging/green/logs",
    ]
    roots.extend(sorted(glob.glob("/home/ec2-user/nexuscrm-deployment/offline/envs/test/*/logs")))
    return roots


def collect_run_scoped_log_actions(keep_runs: int) -> list[CleanupAction]:
    actions: list[CleanupAction] = []
    if keep_runs < 1:
        keep_runs = 1
    for root in run_scoped_log_dir_roots():
        run_dirs = sorted(
            (child for child in Path(root).glob("run-*") if child.is_dir()),
            key=lambda p: p.stat().st_mtime,
            reverse=True,
        )
        for stale in run_dirs[keep_runs:]:
            actions.append(
                CleanupAction(
                    category="run_scoped_log",
                    path=str(stale),
                    size_bytes=path_size(stale),
                    reason=f"run-scoped backend log dir beyond keep-runs={keep_runs}",
                    operation="remove_tree",
                )
            )
    return actions


def collect_log_actions(max_log_bytes: int) -> list[CleanupAction]:
    actions: list[CleanupAction] = []
    for path in managed_log_paths():
        size = file_size(path)
        if size <= max_log_bytes:
            continue
        actions.append(
            CleanupAction(
                category="runtime_log",
                path=path,
                size_bytes=size,
                reason=f"log exceeds {format_bytes(max_log_bytes)}",
                operation="truncate",
            )
        )
    return actions


def runner_roots() -> list[Path]:
    return sorted(Path("/home/ec2-user").glob("actions-runner-*"))


def active_runner_roots() -> set[Path]:
    output = run_text(["ps", "-eo", "args="])
    active: set[Path] = set()
    for root in runner_roots():
        root_text = str(root)
        for line in output.splitlines():
            # Runner.Worker/Plugin processes exist while a job is executing.
            if root_text in line and ("Runner.Worker" in line or "Runner.PluginHost" in line):
                active.add(root)
                break
    return active


def collect_runner_actions(temp_max_age_hours: int, diag_max_age_days: int) -> list[CleanupAction]:
    actions: list[CleanupAction] = []
    now = time.time()
    temp_cutoff = now - (temp_max_age_hours * 3600)
    diag_cutoff = now - (diag_max_age_days * 24 * 3600)
    active_roots = active_runner_roots()

    for root in runner_roots():
        for tarball in root.glob("*.tar.gz"):
            size = path_size(tarball)
            if size > 0:
                actions.append(
                    CleanupAction(
                        category="runner_install_archive",
                        path=str(tarball),
                        size_bytes=size,
                        reason="self-hosted runner install archive is not needed after install",
                        operation="remove_file",
                    )
                )

        diag = root / "_diag"
        if diag.is_dir():
            for item in diag.iterdir():
                if item.is_file() and older_than(item, diag_cutoff):
                    actions.append(
                        CleanupAction(
                            category="runner_diag",
                            path=str(item),
                            size_bytes=path_size(item),
                            reason=f"runner diagnostic file older than {diag_max_age_days} days",
                            operation="remove_file",
                        )
                    )

        if root in active_roots:
            continue

        work = root / "_work"
        if not work.is_dir():
            continue
        candidates = [
            work / "_temp",
            work / "_actions",
            work / "_PipelineMapping",
        ]
        candidates.extend(child for child in work.iterdir() if child.is_dir() and not child.name.startswith("_"))
        for item in candidates:
            if item.exists() and older_than(item, temp_cutoff):
                actions.append(
                    CleanupAction(
                        category="runner_work",
                        path=str(item),
                        size_bytes=path_size(item),
                        reason=f"idle runner work directory older than {temp_max_age_hours} hours",
                        operation="remove_tree",
                    )
                )
    return actions


def collect_managed_temp_actions(max_age_hours: int) -> list[CleanupAction]:
    actions: list[CleanupAction] = []
    cutoff = time.time() - (max_age_hours * 3600)
    paths = [
        Path("/home/ec2-user/nexuscrm-deployment/.tmp/ci-build"),
        Path("/home/ec2-user/nexuscrm-deployment/offline/envs/test/shared/tmp"),
    ]
    for path in paths:
        if path.exists() and older_than(path, cutoff):
            actions.append(
                CleanupAction(
                    category="managed_temp",
                    path=str(path),
                    size_bytes=path_size(path),
                    reason=f"managed host temp directory older than {max_age_hours} hours",
                    operation="remove_tree",
                )
            )
    return actions


def parse_date_dir_name(name: str) -> date | None:
    try:
        return datetime.strptime(name, "%Y-%m-%d").date()
    except ValueError:
        return None


def collect_sf_export_actions(keep_days: int) -> list[CleanupAction]:
    actions: list[CleanupAction] = []
    if keep_days < 1:
        keep_days = 1
    cutoff = date.today() - timedelta(days=keep_days - 1)
    roots = [
        Path("/home/ec2-user/sf-export"),
        Path("/home/ec2-user/sf-export-hourly"),
    ]
    for root in roots:
        if not root.is_dir():
            continue
        for child in root.iterdir():
            if not child.is_dir():
                continue
            child_date = parse_date_dir_name(child.name)
            if child_date is None or child_date >= cutoff:
                continue
            actions.append(
                CleanupAction(
                    category="sf_export",
                    path=str(child),
                    size_bytes=path_size(child),
                    reason=f"Salesforce export date {child.name} older than keep-days={keep_days}",
                    operation="remove_tree",
                )
            )
    return actions


def run_text(args: list[str]) -> str:
    try:
        return subprocess.check_output(args, text=True, stderr=subprocess.STDOUT).strip()
    except (OSError, subprocess.CalledProcessError):
        return ""


def journal_usage_bytes() -> int:
    output = run_text(["journalctl", "--disk-usage"])
    match = re.search(r"take up ([0-9.]+)([KMGTP])", output)
    if not match:
        return 0
    value = float(match.group(1))
    scale = {"K": 1024, "M": 1024**2, "G": 1024**3, "T": 1024**4, "P": 1024**5}[match.group(2)]
    return int(value * scale)


def docker_container_size(container_name: str) -> dict[str, str]:
    output = run_text(["docker", "ps", "-a", "--size", "--filter", f"name=^{container_name}$", "--format", "{{.Names}}\t{{.Size}}\t{{.Status}}"])
    if not output:
        return {"name": container_name, "size": "", "status": "", "found": "false"}
    first = output.splitlines()[0].split("\t")
    while len(first) < 3:
        first.append("")
    return {"name": first[0], "size": first[1], "status": first[2], "found": "true"}


def execute_actions(actions: list[CleanupAction], *, dry_run: bool) -> tuple[int, int]:
    changed_count = 0
    reclaimed_bytes = 0
    for action in actions:
        if dry_run:
            log(f"[dry-run] {action.operation} {action.category}: {action.path} ({format_bytes(action.size_bytes)})")
            reclaimed_bytes += action.size_bytes
            continue
        if action.operation == "truncate":
            with open(action.path, "r+b") as handle:
                handle.truncate(0)
            log(f"truncated {action.category}: {action.path} ({format_bytes(action.size_bytes)})")
            changed_count += 1
            reclaimed_bytes += action.size_bytes
        elif action.operation == "remove_file":
            try:
                Path(action.path).unlink()
            except FileNotFoundError:
                pass
            log(f"removed {action.category}: {action.path} ({format_bytes(action.size_bytes)})")
            changed_count += 1
            reclaimed_bytes += action.size_bytes
        elif action.operation == "remove_tree":
            shutil.rmtree(action.path, ignore_errors=True)
            log(f"removed {action.category}: {action.path} ({format_bytes(action.size_bytes)})")
            changed_count += 1
            reclaimed_bytes += action.size_bytes
    return changed_count, reclaimed_bytes


def docker_cmd(args: list[str]) -> list[str]:
    if run_text(["docker", "info"]):
        return ["docker", *args]
    return ["sudo", "docker", *args]


def recreate_shared_tidb(container_name: str, image: str, *, dry_run: bool) -> dict[str, object]:
    before = docker_container_size(container_name)
    command = [
        "run",
        "-d",
        "--name",
        container_name,
        "--restart",
        "unless-stopped",
        "-p",
        "4000:4000",
        "-p",
        "10080:10080",
        image,
    ]
    if dry_run:
        return {
            "requested": True,
            "changed": False,
            "dry_run": True,
            "before": before,
            "after": before,
            "operation": f"rm -f {container_name}; docker {' '.join(command)}",
        }

    rm_output = run_text(docker_cmd(["rm", "-f", container_name]))
    run_output = run_text(docker_cmd(command))
    after = docker_container_size(container_name)
    return {
        "requested": True,
        "changed": True,
        "dry_run": False,
        "before": before,
        "after": after,
        "rm_output": rm_output,
        "run_output": run_output,
    }


def write_summary(path: str | None, payload: dict[str, object]) -> None:
    if not path:
        return
    out_path = Path(path).expanduser().resolve()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


def main() -> int:
    args = parse_args()
    actions = collect_log_actions(args.max_log_bytes)
    actions.extend(collect_run_scoped_log_actions(args.log_keep_runs))
    actions.extend(collect_runner_actions(args.runner_temp_max_age_hours, args.runner_diag_max_age_days))
    actions.extend(collect_managed_temp_actions(args.managed_temp_max_age_hours))
    actions.extend(collect_sf_export_actions(args.sf_export_keep_days))
    journal_before = journal_usage_bytes()
    changed_count, reclaimed_bytes = execute_actions(actions, dry_run=bool(args.dry_run))

    journal_vacuum_output = ""
    if not args.dry_run:
        journal_vacuum_output = run_text(["journalctl", f"--vacuum-size={args.journal_vacuum_size}"])

    tidb_recreate = {"requested": False}
    if args.recreate_shared_tidb:
        tidb_recreate = recreate_shared_tidb(args.shared_tidb_container_name, args.shared_tidb_image, dry_run=bool(args.dry_run))

    payload = {
        "target": "host",
        "dry_run": bool(args.dry_run),
        "max_log_bytes": args.max_log_bytes,
        "max_log_human": format_bytes(args.max_log_bytes),
        "planned_action_count": len(actions),
        "changed_count": changed_count,
        "reclaimed_bytes": reclaimed_bytes,
        "reclaimed_human": format_bytes(reclaimed_bytes),
        "journal_usage_before_bytes": journal_before,
        "journal_usage_before_human": format_bytes(journal_before),
        "journal_vacuum_size": args.journal_vacuum_size,
        "journal_vacuum_output": journal_vacuum_output,
        "runner_temp_max_age_hours": args.runner_temp_max_age_hours,
        "runner_diag_max_age_days": args.runner_diag_max_age_days,
        "managed_temp_max_age_hours": args.managed_temp_max_age_hours,
        "sf_export_keep_days": args.sf_export_keep_days,
        "log_keep_runs": args.log_keep_runs,
        "shared_tidb_container": docker_container_size(args.shared_tidb_container_name),
        "shared_tidb_recreate": tidb_recreate,
        "actions": [
            {
                "category": action.category,
                "path": action.path,
                "size_bytes": action.size_bytes,
                "size_human": format_bytes(action.size_bytes),
                "reason": action.reason,
                "operation": action.operation,
            }
            for action in actions
        ],
    }
    write_summary(args.summary_json, payload)
    print(json.dumps(payload, indent=2, sort_keys=True))
    return 0


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