#!/usr/bin/env python3
"""
Print Rich tables of .T04 file alignment by 12-char timestamp
(last 12 characters of the basename before .T04) for given directories.
The receiver/source prefix is auto-discovered as everything in the basename
before that 12-character timestamp.

By default, immediate subdirectories of the script’s folder are scanned; pass
``DIR`` paths or set ``--root`` to change that. See ``--help``.
"""
from __future__ import annotations

import argparse
import os
import sys
from collections import defaultdict
from typing import DefaultDict, Iterable, List, Set, Tuple

try:
    from rich import box
    from rich.console import Console
    from rich.markup import escape
    from rich.table import Table
except ModuleNotFoundError as _e:
    _m = _e.name or ""
    if _m != "rich" and not _m.startswith("rich."):
        raise
    _here = os.path.dirname(os.path.abspath(__file__))
    _req = os.path.join(_here, "requirements.txt")
    print(
        "The 'rich' package is not installed for the Python you used to run this script.\n"
        f"  Python:  {sys.executable}\n"
        "  Fix:     install into that same environment, e.g.\n"
        f"            {sys.executable} -m pip install rich\n"
        "           or, from this project directory:\n"
        f"            {sys.executable} -m pip install -r {_req}\n"
        "\n"
        "  (Running plain `pip install` can target a different interpreter than `python3`.)",
        file=sys.stderr,
    )
    raise SystemExit(1) from None

Timestamp = str
Prefix = str
ByTimestamp = DefaultDict[Timestamp, dict[Prefix, str]]


def collect_prefixes(by_ts: ByTimestamp) -> List[Prefix]:
    s = {p for m in by_ts.values() for p in m}
    return sorted(s)


def format_col_header(pfx: Prefix) -> str:
    p = pfx if pfx else "<empty prefix>"
    return escape(p)


def parse_t04_basename(name: str) -> Tuple[Prefix, Timestamp] | None:
    if not name.lower().endswith(".t04"):
        return None
    stem = name[:-4]
    if len(stem) < 12:
        return None
    ts = stem[-12:]
    if not ts.isdigit():
        return None
    return stem[:-12], ts


def scan_dir(dirpath: str) -> ByTimestamp:
    out: ByTimestamp = defaultdict(dict)
    for name in os.listdir(dirpath):
        p = os.path.join(dirpath, name)
        if not os.path.isfile(p):
            continue
        r = parse_t04_basename(name)
        if r is None:
            continue
        prefix, ts = r
        out[ts][prefix] = name
    return out


def scan_dirs(dirpaths: Iterable[str]) -> ByTimestamp:
    merged: ByTimestamp = defaultdict(dict)
    for d in dirpaths:
        one = scan_dir(d)
        for ts, m in one.items():
            for pre, name in m.items():
                merged[ts][pre] = name
    return merged


def discover_subdirectories(root: str) -> List[str]:
    """
    All immediate subdirectories of ``root`` (not recursive), sorted by name.
    Skips names starting with ".".
    """
    if not os.path.isdir(root):
        return []
    out: List[str] = []
    for name in sorted(os.listdir(root)):
        if name.startswith("."):
            continue
        p = os.path.join(root, name)
        if os.path.isdir(p):
            out.append(os.path.abspath(p))
    return out


def where_present_key(
    dirpaths: List[str],
) -> DefaultDict[Tuple[Prefix, Timestamp], Set[str]]:
    """(prefix, timestamp) -> which folder basenames have that file."""
    out: DefaultDict[Tuple[Prefix, Timestamp], Set[str]] = defaultdict(set)
    for d in dirpaths:
        b = os.path.basename(d)
        for name in os.listdir(d):
            p = os.path.join(d, name)
            if not os.path.isfile(p):
                continue
            r = parse_t04_basename(name)
            if r is None:
                continue
            prefix, ts = r
            out[(prefix, ts)].add(b)
    return out


def _stdout_supports_ansi() -> bool:
    return bool(sys.stdout.isatty() and "NO_COLOR" not in os.environ)


def print_table(
    by_ts: ByTimestamp,
    title: str | None = None,
    show_all_prefixes: bool = False,
) -> None:
    prefixes = collect_prefixes(by_ts)
    n_p = len(prefixes)
    if not by_ts or not prefixes:
        e = Console(
            file=sys.stderr,
            no_color=not (sys.stderr.isatty() and "NO_COLOR" not in os.environ),
        )
        e.print(escape("*(No matching .T04 files.)*"), style="dim")
        return

    use_ansi = _stdout_supports_ansi()
    box_style: box.Box = box.ROUNDED if use_ansi else box.ASCII
    console = Console(
        file=sys.stdout,
        no_color=not use_ansi,
        soft_wrap=True,
    )

    table = Table(
        title=escape(title) if title else None,
        title_justify="left",
        title_style="bold",
        show_header=True,
        show_lines=False,
        header_style="bold",
        box=box_style,
        # Same (vertical, horizontal) padding in header and body; vertical=middle in columns
        padding=(0, 1),
        highlight=False,
        show_edge=True,
    )
    table.add_column(escape("Timestamp (12 chr)"), justify="left", vertical="middle")
    for pfx in prefixes:
        table.add_column(format_col_header(pfx), justify="center", vertical="middle")
    if show_all_prefixes:
        table.add_column(escape(f"All {n_p}?" if n_p else "All?"), justify="center", vertical="middle")

    for ts in sorted(by_ts):
        n = 0
        cells: List[str] = [str(ts)]
        for pfx in prefixes:
            y = pfx in by_ts[ts]
            cells.append("yes" if y else "—")
            if y:
                n += 1
        if show_all_prefixes:
            alln = "yes" if n == n_p and n_p else f"no {n}/{n_p}"
            cells.append(alln)
        table.add_row(*cells)

    n_ts = len(by_ts)
    full = sum(1 for t in by_ts if all(p in by_ts[t] for p in prefixes))
    n_word = f"{n_p} prefix" + ("es" if n_p != 1 else "")

    console.print()
    console.print(table)
    print(
        f"\nCount: {n_ts} timestamps, {full} with all {n_word}.\n",
        file=sys.stderr,
        flush=True,
    )


def main() -> int:
    p = argparse.ArgumentParser(
        description="Print .T04 timestamp alignment table(s) (12-char stamp before .T04); "
        "device columns are the basename prefix before that timestamp, auto-discovered.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument(
        "dirs",
        nargs="*",
        metavar="DIR",
        help="Directories to scan. If omitted, every immediate subdirectory of "
        "--root is used (non-recursive, sorted by name; entries starting with “.” are skipped).",
    )
    p.add_argument(
        "--root",
        type=str,
        default='.',
        help="When DIR arguments are omitted, this directory is scanned for subfolders to use.",
    )
    p.add_argument(
        "--per-folder",
        action="store_true",
        help="Also print a separate table for each directory.",
    )
    p.add_argument(
        "--all-prefixes-column",
        action="store_true",
        dest="all_prefixes_column",
        help="Add a column showing when every auto-discovered prefix is present. ",
    )
    p.add_argument(
        "--duplicates",
        action="store_true",
        help="Report (prefix, timestamp) keys that appear in more than one folder.",
    )
    args = p.parse_args()
    root = os.path.abspath(os.path.expanduser(args.root))
    if args.dirs:
        dirs = [os.path.abspath(os.path.expanduser(d)) for d in args.dirs]
    else:
        dirs = discover_subdirectories(root)
        if not dirs:
            print(
                f"No subdirectories found under: {root}",
                file=sys.stderr,
            )
            return 1
    for d in dirs:
        if not os.path.isdir(d):
            print(f"Not a directory: {d}", file=sys.stderr)
            return 1

    combined = scan_dirs(dirs)
    print_table(
        combined,
        title="Combined: all unique timestamps",
        show_all_prefixes=args.all_prefixes_column,
    )

    if args.per_folder:
        for d in dirs:
            one = scan_dir(d)
            print_table(
                one,
                title=f"Folder `{os.path.basename(d)}` only",
                show_all_prefixes=args.all_prefixes_column,
            )

    if args.duplicates:
        loc = where_present_key(dirs)
        both = [k for k, s in loc.items() if len(s) > 1]
        print("## Same prefix+timestamp in more than one folder", file=sys.stderr)
        if not both:
            print("Count: 0", file=sys.stderr)
        else:
            print(f"Count: {len(both)}", file=sys.stderr)
            for key in sorted(both):
                print(" ", key, "->", loc[key], file=sys.stderr)
    return 0


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