from __future__ import annotations

import argparse
import os
import shutil
import uuid
from collections import Counter
from collections.abc import Iterable, Sequence
from dataclasses import dataclass
from pathlib import Path

from copienator import (
    CliError,
    EvaluationWorkspace,
    ExitCode,
    configuration,
    evaluation_parser,
    execute,
    workspace_from_args,
)


IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp"}
STATEMENT_ROOT_FILES = {"enonce.tex", "correction.tex", "labels", "label_groups"}
STATEMENT_TEXT_DIRECTORIES = {"Text", "Sol", "Persp", "Cache", "Tmp"}


@dataclass(frozen=True, slots=True)
class CleanupPlan:
    kept_files: tuple[Path, ...]
    deleted_files: tuple[Path, ...]
    deleted_directories: tuple[Path, ...]
    bytes_to_delete: int


def _workspace_entries(root: Path) -> tuple[list[Path], list[Path]]:
    """List files/symlinks and real directories without following symlinks."""
    files: list[Path] = []
    directories: list[Path] = []
    for current, directory_names, file_names in os.walk(
        root, topdown=True, followlinks=False
    ):
        current_path = Path(current)
        traversable: list[str] = []
        for name in directory_names:
            path = current_path / name
            if path.is_symlink():
                files.append(path)
            else:
                directories.append(path)
                traversable.append(name)
        directory_names[:] = traversable
        files.extend(current_path / name for name in file_names)
    return files, directories


def _return_artifacts(workspace: EvaluationWorkspace) -> list[Path]:
    return_dir = workspace.return_dir
    if not return_dir.is_dir() or return_dir.is_symlink():
        raise CliError(
            f"Return directory not found or invalid: {return_dir}",
            ExitCode.INVALID_WORKSPACE,
        )

    student_directories = sorted(
        (
            path
            for path in return_dir.iterdir()
            if path.is_dir() and not path.is_symlink()
        ),
        key=lambda path: path.name.casefold(),
    )
    if not student_directories:
        raise CliError(
            f"No student directories found in {return_dir}",
            ExitCode.INVALID_WORKSPACE,
        )

    artifacts: list[Path] = []
    incomplete: list[str] = []
    for directory in student_directories:
        files = [path for path in directory.rglob("*") if path.is_file()]
        images = [
            path for path in files if path.suffix.casefold() in IMAGE_SUFFIXES
        ]
        scores = [path for path in files if path.name.casefold() == "score.json"]
        pdfs = [path for path in files if path.suffix.casefold() == ".pdf"]
        if (configuration.RETURN_JPEG_ENABLED and not images) or not scores:
            missing = []
            if configuration.RETURN_JPEG_ENABLED and not images:
                missing.append("image")
            if not scores:
                missing.append("score.json")
            incomplete.append(f"{directory.name} ({', '.join(missing)})")
        artifacts.extend(images)
        artifacts.extend(scores)
        artifacts.extend(pdfs)
        artifacts.extend(path for path in files if path.name.casefold() == "info.json")

    if incomplete:
        details = "\n".join(f" - {item}" for item in incomplete)
        raise CliError(
            "A Rendre is incomplete; cleanup was refused:\n" + details,
            ExitCode.INVALID_WORKSPACE,
        )
    return artifacts


def _is_statement_text_file(workspace: EvaluationWorkspace, path: Path) -> bool:
    relative = path.relative_to(workspace.root)
    if len(relative.parts) == 1:
        return relative.name in STATEMENT_ROOT_FILES
    top_level = relative.parts[0]
    if top_level in STATEMENT_TEXT_DIRECTORIES:
        return path.suffix.casefold() in {"", ".json", ".tex", ".txt"}
    if top_level in {"Text2", "Sol2"}:
        return path.suffix.casefold() == ".tex"
    return False


def _is_log_file(workspace: EvaluationWorkspace, path: Path) -> bool:
    if path.is_relative_to(workspace.logs_dir):
        return True
    relative = path.relative_to(workspace.root)
    return len(relative.parts) == 1 and (
        path.suffix.casefold() == ".log"
        or path.name.casefold().endswith("_log")
    )


def build_cleanup_plan(workspace: EvaluationWorkspace) -> CleanupPlan:
    processed_copies = sorted(
        (
            path
            for path in workspace.copies_dir.glob("*.pdf")
            if path.is_file()
        ),
        key=lambda path: path.name.casefold(),
    )
    if not processed_copies:
        raise CliError(
            f"No processed PDF copies found in {workspace.copies_dir}",
            ExitCode.INVALID_WORKSPACE,
        )
    if not workspace.correction_file.is_file():
        raise CliError(
            f"Correction result not found: {workspace.correction_file}",
            ExitCode.INVALID_WORKSPACE,
        )

    files, directories = _workspace_entries(workspace.root)

    kept_files = set(processed_copies)
    kept_files.add(workspace.correction_file)
    kept_files.update(_return_artifacts(workspace))
    kept_files.update(
        path
        for path in files
        if _is_statement_text_file(workspace, path)
        or _is_log_file(workspace, path)
    )

    kept_directories = {workspace.root}
    for path in kept_files:
        kept_directories.update(
            parent
            for parent in path.parents
            if parent == workspace.root or workspace.root in parent.parents
        )

    deleted_files = tuple(sorted(set(files) - kept_files, key=str))
    deleted_directories = tuple(
        sorted(
            set(directories) - kept_directories,
            key=lambda path: (len(path.parts), str(path)),
            reverse=True,
        )
    )
    bytes_to_delete = sum(
        path.stat(follow_symlinks=False).st_size
        for path in deleted_files
        if path.exists() or path.is_symlink()
    )
    return CleanupPlan(
        kept_files=tuple(sorted(kept_files, key=str)),
        deleted_files=deleted_files,
        deleted_directories=deleted_directories,
        bytes_to_delete=bytes_to_delete,
    )


def _human_size(size: int) -> str:
    value = float(size)
    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
        if value < 1024 or unit == "TiB":
            return f"{value:.1f} {unit}"
        value /= 1024
    return f"{size} B"


def _top_level_counts(
    workspace: EvaluationWorkspace, paths: Iterable[Path]
) -> Counter[str]:
    counts: Counter[str] = Counter()
    for path in paths:
        relative = path.relative_to(workspace.root)
        counts[relative.parts[0]] += 1
    return counts


def print_plan(workspace: EvaluationWorkspace, plan: CleanupPlan, *, verbose: bool) -> None:
    processed_count = sum(path.parent == workspace.copies_dir for path in plan.kept_files)
    return_count = sum(
        path.is_relative_to(workspace.return_dir) for path in plan.kept_files
    )
    statement_count = sum(
        _is_statement_text_file(workspace, path) for path in plan.kept_files
    )
    log_count = sum(_is_log_file(workspace, path) for path in plan.kept_files)
    print(f"Evaluation: {workspace.root}")
    print("Will keep:")
    print(f" - {processed_count} processed PDF copies in Copies")
    print(f" - {statement_count} textual statement files")
    print(" - correction.json")
    print(f" - {log_count} log files")
    print(f" - {return_count} image/PDF/score artifacts in A Rendre")
    print(
        f"Will delete {len(plan.deleted_files)} files and "
        f"{len(plan.deleted_directories)} directories "
        f"({_human_size(plan.bytes_to_delete)} in file entries)."
    )
    counts = _top_level_counts(workspace, plan.deleted_files)
    if counts:
        print("Files removed by top-level location:")
        for name, count in sorted(counts.items(), key=lambda item: item[0].casefold()):
            print(f" - {name}: {count}")
    if verbose:
        print("Deletion list:")
        for path in plan.deleted_files:
            print(f" - {path.relative_to(workspace.root)}")


def _materialize_return_links(workspace: EvaluationWorkspace, paths: Iterable[Path]) -> None:
    for path in paths:
        if not path.is_relative_to(workspace.return_dir) or not path.is_symlink():
            continue
        try:
            source = path.resolve(strict=True)
        except OSError as exc:
            raise CliError(f"Broken return link {path}: {exc}") from exc
        if not source.is_file():
            raise CliError(f"Return link does not target a file: {path} -> {source}")
        temporary = path.with_name(f".{path.name}.materialize-{uuid.uuid4().hex}.tmp")
        try:
            shutil.copy2(source, temporary)
            temporary.replace(path)
        finally:
            temporary.unlink(missing_ok=True)
        print(f"Materialized: {path.relative_to(workspace.root)}")


def apply_cleanup(workspace: EvaluationWorkspace, plan: CleanupPlan) -> None:
    _materialize_return_links(workspace, plan.kept_files)
    for path in plan.deleted_files:
        path.unlink(missing_ok=True)
    for path in plan.deleted_directories:
        try:
            path.rmdir()
        except FileNotFoundError:
            pass
    print(
        f"Cleanup complete: deleted {len(plan.deleted_files)} files and "
        f"{len(plan.deleted_directories)} directories."
    )


def run(
    workspace: EvaluationWorkspace,
    *,
    dry_run: bool = False,
    assume_yes: bool = False,
    verbose: bool = False,
) -> ExitCode:
    plan = build_cleanup_plan(workspace)
    print_plan(workspace, plan, verbose=verbose)
    if dry_run:
        print("Dry run: nothing was deleted.")
        return ExitCode.SUCCESS
    if not assume_yes:
        expected = workspace.name
        answer = input(
            f"This cannot be undone. Type {expected!r} to confirm cleanup: "
        ).strip()
        if answer != expected:
            print("Cleanup cancelled; nothing was deleted.")
            return ExitCode.SUCCESS
    apply_cleanup(workspace, plan)
    return ExitCode.SUCCESS


def build_parser() -> argparse.ArgumentParser:
    parser = evaluation_parser(
        "Archive an evaluation by deleting regenerable and intermediate files."
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be kept and deleted without changing anything",
    )
    parser.add_argument(
        "--yes",
        action="store_true",
        help="Skip the interactive evaluation-name confirmation",
    )
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()

    def handle(args: argparse.Namespace) -> ExitCode:
        return run(
            workspace_from_args(args),
            dry_run=args.dry_run,
            assume_yes=args.yes,
            verbose=args.verbose,
        )

    return execute(parser, argv, handle)


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