#!/usr/bin/env python3
"""
Folder Compressor using 7-Zip
- Select a single folder to compress, OR
- Compress all folders one by one automatically
Uses 7-Zip store mode (no compression) for maximum speed.
Real-time progress bar shown during each compression job.
"""

import os
import re
import sys
import shutil
import subprocess
import time
from pathlib import Path
from dataclasses import dataclass


# ── Terminal color helpers ────────────────────────────────────────────────────

SUPPORTS_COLOR = (
    sys.stdout.isatty()
    and (sys.platform != "win32" or os.environ.get("WT_SESSION") or os.environ.get("TERM"))
)

def _c(code: str, text: str) -> str:
    return f"\033[{code}m{text}\033[0m" if SUPPORTS_COLOR else text

def green(t):   return _c("32", t)
def yellow(t):  return _c("33", t)
def cyan(t):    return _c("36", t)
def bold(t):    return _c("1",  t)
def dim(t):     return _c("2",  t)
def red(t):     return _c("31", t)
def magenta(t): return _c("35", t)


# ── Size / time formatters ────────────────────────────────────────────────────

def fmt_size(b: float) -> str:
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if b < 1024:
            return f"{b:.1f} {unit}"
        b /= 1024
    return f"{b:.1f} PB"


def fmt_time(seconds: float) -> str:
    s = int(seconds)
    if s < 60:
        return f"{s}s"
    m, s = divmod(s, 60)
    if m < 60:
        return f"{m}m {s:02d}s"
    h, m = divmod(m, 60)
    return f"{h}h {m:02d}m {s:02d}s"


# ── Folder size calculator ────────────────────────────────────────────────────

def folder_size(path: Path) -> int:
    total = 0
    for root, _, files in os.walk(path):
        for f in files:
            try:
                total += os.path.getsize(os.path.join(root, f))
            except OSError:
                pass
    return total


# ── 7-Zip locator ─────────────────────────────────────────────────────────────

def find_7zip() -> str | None:
    for name in ("7z", "7za", "7zz", "7zip"):
        path = shutil.which(name)
        if path:
            return path
    if sys.platform == "win32":
        for p in (
            r"C:\Program Files\7-Zip\7z.exe",
            r"C:\Program Files (x86)\7-Zip\7z.exe",
        ):
            if os.path.isfile(p):
                return p
    return None


# ── Folder listing ────────────────────────────────────────────────────────────

def list_folders(base_dir: str) -> list[Path]:
    base = Path(base_dir)
    return sorted(
        (e for e in base.iterdir() if e.is_dir()),
        key=lambda p: p.name.lower(),
    )


# ── Mode + folder selection menu ─────────────────────────────────────────────

def prompt_mode_and_selection(folders: list[Path]) -> list[Path]:
    """
    Print the folder list then ask the user what to do:
      A  — compress ALL folders sequentially
      1…N — compress a single specific folder
    Returns the list of folders to compress (one or all).
    """
    print("\nAvailable folders:\n")
    for i, folder in enumerate(folders, start=1):
        try:
            sz = folder_size(folder)
            size_label = dim(f"  ({fmt_size(sz)})")
        except PermissionError:
            size_label = dim("  (no access)")
        print(f"  {dim(f'[{i:>3}]')}  {folder.name}{size_label}")

    print()
    print(f"  {dim('[  A]')}  {bold('Compress ALL folders one by one')}")
    print()

    while True:
        try:
            raw = input("Enter a folder number or  A  to compress all: ").strip().upper()

            if raw == "A":
                print(f"\n  {green('✔')}  Mode: {bold('Compress all folders sequentially')}")
                return list(folders)

            idx = int(raw)
            if 1 <= idx <= len(folders):
                chosen = folders[idx - 1]
                print(f"\n  {green('✔')}  Mode: {bold('Single folder')}  →  {chosen.name}")
                return [chosen]

            print(f"  Please enter a number between 1 and {len(folders)}, or A.")

        except ValueError:
            print("  Invalid input — enter a number or A.")
        except (EOFError, KeyboardInterrupt):
            print("\nAborted.")
            sys.exit(0)


# ── Progress bar renderer ─────────────────────────────────────────────────────

BAR_WIDTH = 36

_PROGRESS_RE = re.compile(r"(\d{1,3})%(?:\s+(\d+)\s*\+?\s*(.*))?")

# Split on backspaces (\x08, Linux/macOS) or carriage returns (\r, Windows)
_SPLIT_RE = re.compile(rb"[\x08\r]+")


def _move_up(n: int):
    if SUPPORTS_COLOR:
        sys.stdout.write(f"\033[{n}A")


def _clear_line():
    if SUPPORTS_COLOR:
        sys.stdout.write("\033[2K\r")
    else:
        sys.stdout.write("\r" + " " * 78 + "\r")


def _redraw(pct: int, elapsed: float, files_done: int,
            current_file: str, speed_bps: float):
    filled  = int(BAR_WIDTH * pct / 100)
    bar     = "█" * filled + "░" * (BAR_WIDTH - filled)
    color   = green if pct == 100 else cyan
    speed_s = f"{fmt_size(speed_bps)}/s" if speed_bps > 0 else "─"

    _move_up(2)

    _clear_line()
    sys.stdout.write(
        f"  {color(bar)}  {bold(f'{pct:>3}%')}"
        f"  {dim('elapsed')} {dim(fmt_time(elapsed))}"
        f"  {yellow(speed_s)}"
        f"  {dim(f'files: {files_done}')}\n"
    )

    _clear_line()
    label = current_file.strip() or "starting…"
    if len(label) > 60:
        label = "…" + label[-59:]
    sys.stdout.write(f"  {dim('↳ ' + label)}\n")

    sys.stdout.flush()


# ── Result record ─────────────────────────────────────────────────────────────

@dataclass
class JobResult:
    folder:      Path
    archive:     Path | None
    src_size:    int
    arc_size:    int
    elapsed:     float
    avg_speed:   float
    success:     bool
    error:       str = ""


# ── Single compression job ────────────────────────────────────────────────────

def compress_folder(seven_zip: str, folder: Path, output_dir: Path,
                    job_label: str = "") -> JobResult:
    archive_path = output_dir / (folder.name + ".7z")

    if job_label:
        print(f"\n{job_label}")

    print(f"\n  {dim('Calculating source size…')} ", end="", flush=True)
    src_size = folder_size(folder)
    print(fmt_size(src_size))

    cmd = [
        seven_zip,
        "a",
        "-t7z",
        "-mx=0",
        "-ms=off",
        "-bsp1",
        str(archive_path),
        str(folder),
    ]

    print(f"\n  {dim('Command: ' + ' '.join(cmd))}")
    print()
    print("  " + "─" * 64)
    print()
    print()

    start_time   = time.monotonic()
    files_done   = 0
    current_file = ""
    speed_bps    = 0.0
    last_bytes   = 0.0
    speed_ts     = start_time

    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )

    buf = b""

    while True:
        chunk = proc.stdout.read(256)
        if not chunk:
            break

        buf    += chunk
        elapsed = time.monotonic() - start_time
        segments = _SPLIT_RE.split(buf)
        buf = segments[-1]

        for seg in segments[:-1]:
            text = seg.decode("utf-8", errors="replace").strip()
            if not text:
                continue

            m = _PROGRESS_RE.search(text)
            if m:
                pct        = min(int(m.group(1)), 100)
                files_done = int(m.group(2)) if m.group(2) else files_done
                fname      = (m.group(3) or "").strip()
                if fname:
                    current_file = fname

                if src_size > 0 and pct > 0:
                    bytes_done = src_size * pct / 100
                    now        = time.monotonic()
                    dt         = now - speed_ts
                    if dt >= 0.4:
                        speed_bps  = (bytes_done - last_bytes) / dt
                        last_bytes = bytes_done
                        speed_ts   = now

                _redraw(pct, elapsed, files_done, current_file, speed_bps)

    proc.wait()
    elapsed_total = time.monotonic() - start_time

    if proc.returncode == 0:
        _redraw(100, elapsed_total, files_done, current_file, speed_bps)

    print()
    print("  " + "─" * 64)

    if proc.returncode == 0:
        arc_size  = archive_path.stat().st_size
        avg_speed = src_size / elapsed_total if elapsed_total > 0 else 0

        print(f"\n  {green('✅  Archive created successfully!')}")
        print(f"     {bold('Path')}     : {archive_path}")
        print(f"     {bold('Source')}   : {fmt_size(src_size)}")
        print(f"     {bold('Archive')}  : {fmt_size(arc_size)}")
        print(f"     {bold('Time')}     : {fmt_time(elapsed_total)}")
        print(f"     {bold('Avg speed')}: {fmt_size(avg_speed)}/s")

        return JobResult(
            folder=folder, archive=archive_path,
            src_size=src_size, arc_size=arc_size,
            elapsed=elapsed_total, avg_speed=avg_speed,
            success=True,
        )
    else:
        msg = f"7-Zip exited with code {proc.returncode}"
        print(f"\n  {red(f'❌  {msg}')}")
        return JobResult(
            folder=folder, archive=None,
            src_size=src_size, arc_size=0,
            elapsed=elapsed_total, avg_speed=0,
            success=False, error=msg,
        )


# ── Batch summary ─────────────────────────────────────────────────────────────

def print_batch_summary(results: list[JobResult]):
    succeeded = [r for r in results if r.success]
    failed    = [r for r in results if not r.success]
    total_src = sum(r.src_size for r in results)
    total_arc = sum(r.arc_size for r in results)
    total_t   = sum(r.elapsed  for r in results)

    width = 66
    print()
    print("  " + "═" * width)
    print(f"  {bold('BATCH SUMMARY')}  —  {len(results)} folder(s) processed")
    print("  " + "═" * width)

    for r in results:
        icon  = green("✅") if r.success else red("❌")
        label = r.folder.name
        if len(label) > 38:
            label = label[:35] + "…"
        if r.success:
            detail = f"{fmt_size(r.src_size)}  →  {fmt_time(r.elapsed)}  @  {fmt_size(r.avg_speed)}/s"
        else:
            detail = red(r.error)
        print(f"  {icon}  {label:<38}  {dim(detail)}")

    print("  " + "─" * width)
    print(f"  {bold('Total source size')} : {fmt_size(total_src)}")
    print(f"  {bold('Total archive size')}: {fmt_size(total_arc)}")
    print(f"  {bold('Total time')}        : {fmt_time(total_t)}")
    if total_t > 0:
        avg = total_src / total_t
        print(f"  {bold('Overall avg speed')} : {fmt_size(avg)}/s")
    if failed:
        print()
        print(f"  {red(f'⚠  {len(failed)} job(s) failed:')}")
        for r in failed:
            print(f"     • {r.folder.name}  —  {r.error}")
    print("  " + "═" * width)
    print()


# ── Entry point ───────────────────────────────────────────────────────────────

def main():
    scan_dir = os.path.abspath(sys.argv[1] if len(sys.argv) > 1 else os.getcwd())

    if not os.path.isdir(scan_dir):
        print(f"Error: '{scan_dir}' is not a valid directory.")
        sys.exit(1)

    print(f"\n{bold('📦  Folder Compressor')}")
    print(f"    Scanning: {cyan(scan_dir)}\n")

    seven_zip = find_7zip()
    if not seven_zip:
        print(
            f"\n{red('❌  7-Zip not found.')}\n"
            "    Install it and make sure '7z' (or '7za') is on your PATH.\n"
            "    • Linux  : sudo apt install p7zip-full\n"
            "    • macOS  : brew install p7zip\n"
            "    • Windows: https://www.7-zip.org/download.html"
        )
        sys.exit(1)

    print(f"  {green('✔')}  7-Zip: {dim(seven_zip)}")

    folders = list_folders(scan_dir)
    if not folders:
        print(f"\nNo sub-folders found in '{scan_dir}'.")
        sys.exit(0)

    targets = prompt_mode_and_selection(folders)

    batch = len(targets) > 1
    results: list[JobResult] = []

    for i, folder in enumerate(targets, start=1):
        label = ""
        if batch:
            label = (
                f"\n  {magenta(f'[{i}/{len(targets)}]')}  "
                f"{bold(folder.name)}"
                f"  {dim('─' * max(0, 44 - len(folder.name)))}"
            )

        result = compress_folder(seven_zip, folder, folder.parent, job_label=label)
        results.append(result)

        # In batch mode, allow Ctrl+C between jobs to abort remaining work
        if batch and i < len(targets):
            try:
                time.sleep(0)  # yield to allow KeyboardInterrupt to surface
            except KeyboardInterrupt:
                print(f"\n  {yellow('⚠  Batch interrupted by user.')}")
                break

    if batch:
        print_batch_summary(results)
    else:
        print()


if __name__ == "__main__":
    main()