#!/usr/bin/env python3

from __future__ import annotations

import argparse
import re
from collections import defaultdict
from datetime import datetime, timedelta
from pathlib import Path

MAX_FILES_PER_GROUP = 10
MAX_AGE_DAYS = 30
MAX_SIZE_BYTES = 10 * 1024 * 1024  # 10 MB

TIMESTAMP_RE = re.compile(r"^(.*?)_(\d{8}_\d{6})\.log$")
LOG_SUFFIX_RE = re.compile(r"^(.*?)\.log(?:_\d+)?$")


def get_log_group(log_file: Path) -> str:
    """
    Group examples:

    ffmpeg_handling.log                    -> ffmpeg_handling
    ffmpeg_handling_20261027_193908.log    -> ffmpeg_handling
    sonarr_translation_tool_*.log         -> sonarr_translation_tool
    log_backup_102214.log                 -> log_backup
    """

    stem = log_file.stem

    match = TIMESTAMP_RE.match(log_file.name)
    if match:
        return match.group(1)

    parts = stem.split("_")

    while parts and (
        re.fullmatch(r"\d{6}", parts[-1])
        or re.fullmatch(r"\d{8}", parts[-1])
        or re.fullmatch(r"\d{8}_\d{6}", parts[-1])
    ):
        parts.pop()

    if parts:
        return "_".join(parts)

    match = LOG_SUFFIX_RE.match(log_file.name)
    if match:
        return match.group(1)

    return stem


def truncate_first_half(file_path: Path) -> None:
    """Remove the first half of a file, keeping the newest half."""

    size = file_path.stat().st_size

    if size <= MAX_SIZE_BYTES:
        return

    with file_path.open("rb") as f:
        f.seek(size // 2)
        data = f.read()

    newline_pos = data.find(b"\n")
    if newline_pos != -1:
        data = data[newline_pos + 1 :]

    with file_path.open("wb") as f:
        f.write(data)

    print(f"TRUNCATED  {file_path.name}")


def remove_old_logs(log_dir: Path) -> None:
    cutoff = datetime.now() - timedelta(days=MAX_AGE_DAYS)

    for file in log_dir.glob("*.log"):
        if datetime.fromtimestamp(file.stat().st_mtime) < cutoff:
            print(f"REMOVED OLD {file.name}")
            file.unlink(missing_ok=True)


def keep_latest_logs(log_dir: Path) -> None:
    groups: dict[str, list[Path]] = defaultdict(list)

    for file in log_dir.glob("*.log"):
        groups[get_log_group(file)].append(file)

    for group_name, files in groups.items():
        files.sort(
            key=lambda f: f.stat().st_mtime,
            reverse=True,
        )

        for old_file in files[MAX_FILES_PER_GROUP:]:
            print(f"REMOVED EXCESS {old_file.name}")
            old_file.unlink(missing_ok=True)


def shrink_large_logs(log_dir: Path) -> None:
    for file in log_dir.glob("*.log"):
        truncate_first_half(file)


def manage_logs(log_dir: Path) -> None:
    remove_old_logs(log_dir)
    keep_latest_logs(log_dir)
    shrink_large_logs(log_dir)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "log_dir",
        nargs="?",
        default=Path("/home/lucas/logs"),
        type=Path,
        help="Directory containing log files",
        )
    args = parser.parse_args()
    manage_logs(args.log_dir.expanduser().resolve())


if __name__ == "__main__":
    main()