main.py

"""Smart Codebase Bundler — on-demand CLI entry point (SRS v8.0 pipeline)."""

from __future__ import annotations

import os
import sys
from pathlib import Path

from colorama import Fore, Style, init as colorama_init

from bundler.cache import (
    BundleCache,
    load_cache,
    md5_file,
    module_needs_rebuild,
    save_cache,
)
from bundler.cli import run
from bundler.scanner import normalize_path, scan
from bundler.writer import (
    BundleWriter,
    GENERAL_MODULE_ID,
    group_files_by_module,
    pack_module_units,
    prepare_source_unit,
    relative_key,
)

colorama_init(autoreset=True)

_active_writer: BundleWriter | None = None


def request_cancel() -> None:
    """Signal the active writer to stop (does not delete TEMP yet)."""
    writer = _active_writer
    if writer is not None:
        writer.request_cancel()


def cleanup() -> None:
    """FR-1.2: cancel then remove TEMP_DIR artifacts on SIGINT/SIGTERM/GUI close."""
    global _active_writer
    writer = _active_writer
    _active_writer = None
    if writer is not None:
        writer.request_cancel()
        writer.cleanup_temp()


def _default_output_dir(project_root: Path) -> Path:
    env = os.environ.get("SMART_BUNDLER_OUTPUT")
    if env:
        return normalize_path(env)
    return normalize_path(project_root / "bundles")


def _default_project_root() -> Path:
    env = os.environ.get("SMART_BUNDLER_ROOT")
    if env:
        return normalize_path(env)
    return normalize_path(Path.cwd())


def _print(msg: str) -> None:
    sys.stdout.write(msg + "\n")


def _print_dry_run(modules_report: list[tuple[str, int, int, bool]]) -> None:
    """Report modules scheduled for bundling and estimated tokens (FR-1.1)."""
    _print(f"{Fore.CYAN}Dry-run - no files will be written{Style.RESET_ALL}")
    total = 0
    for module_id, file_count, tokens, rebuild in modules_report:
        flag = "REBUILD" if rebuild else "skip (cache)"
        _print(
            f"  [{flag}] bundle_{module_id}: {file_count} files, ~{tokens} tokens"
        )
        if rebuild:
            total += tokens
    _print(f"Estimated tokens to rebuild: ~{total}")


def run_pipeline(*, force: bool = False, dry_run: bool = False) -> int:
    global _active_writer

    project_root = _default_project_root()
    output_dir = _default_output_dir(project_root)
    if not dry_run:
        output_dir.mkdir(parents=True, exist_ok=True)

    # Stage 2 — scan (never re-ingest OUTPUT_DIR)
    scanned = []
    out_norm = normalize_path(output_dir)
    for path in scan(project_root):
        try:
            normalize_path(path).relative_to(out_norm)
        except ValueError:
            scanned.append(path)
    grouped = group_files_by_module(project_root, scanned)
    cache = load_cache(project_root)

    writer = BundleWriter(output_dir)
    _active_writer = writer

    try:
        modules_to_build = []
        dry_rows: list[tuple[str, int, int, bool]] = []
        active_names: set[str] = set()
        updated_cache = BundleCache(
            project_root=str(project_root),
            files=dict(cache.files),
            modules={},
        )

        module_ids = sorted(
            grouped.keys(),
            key=lambda m: (0, m) if m == GENERAL_MODULE_ID else (1, m.lower()),
        )

        for module_id in module_ids:
            paths = sorted(
                grouped[module_id],
                key=lambda p: relative_key(p, project_root),
            )
            keys = [relative_key(p, project_root) for p in paths]

            # Stage 3 — cheap MD5 probe for desync / change detection
            file_md5s: dict[str, str] = {}
            for path, key in zip(paths, keys):
                digest = md5_file(path)
                file_md5s[key] = digest or ""

            rebuild = module_needs_rebuild(
                module_id,
                keys,
                file_md5s,
                cache,
                output_dir,
                force=force,
            )

            if not rebuild:
                token_sum = 0
                for key in keys:
                    entry = cache.file_entry(key)
                    if entry:
                        token_sum += entry.tokens
                        updated_cache.set_file(key, entry.md5, entry.tokens)
                prev = cache.modules.get(module_id)
                if prev:
                    updated_cache.set_module(module_id, prev.parts, prev.files)
                    active_names.update(prev.parts)
                dry_rows.append((module_id, len(keys), token_sum, False))
                continue

            # Stage 4 — full read + hybrid tokens only for modules that rebuild
            units = []
            token_sum = 0
            for path, key in zip(paths, keys):
                cached = cache.file_entry(key)
                disk_md5 = file_md5s.get(key) or ""
                unit = prepare_source_unit(
                    path,
                    root=project_root,
                    rel_key=key,
                    cached_tokens=cached.tokens if cached else None,
                    cached_md5=cached.md5 if cached else None,
                    disk_md5=disk_md5,
                )
                units.append(unit)
                token_sum += unit.cache_tokens
                # Persist disk MD5 + exact tiktoken count (FR-1.5).
                updated_cache.set_file(
                    key,
                    disk_md5 or unit.md5,
                    unit.cache_tokens,
                )

            dry_rows.append((module_id, len(units), token_sum, True))

            bundle = pack_module_units(module_id, units)
            if not bundle.parts:
                continue
            modules_to_build.append(bundle)
            updated_cache.set_module(
                module_id,
                bundle.filenames,
                bundle.source_files,
            )
            active_names.update(bundle.filenames)

        if dry_run:
            _print_dry_run(dry_rows)
            return 0

        for mid, mod in updated_cache.modules.items():
            active_names.update(mod.parts)

        # FR-4.5 — full active module map (rebuilt + cache-hit).
        rebuilt_by_id = {m.module_id: m for m in modules_to_build}
        manifest_modules: dict = {}
        total_tokens = 0
        for mid, mod in updated_cache.modules.items():
            if mid in rebuilt_by_id:
                bundle = rebuilt_by_id[mid]
                entry = {
                    "parts": bundle.filenames,
                    "part_count": len(bundle.parts),
                    "estimated_tokens": bundle.total_tokens,
                    "files": bundle.source_files,
                }
            else:
                file_token_sum = 0
                for key in mod.files:
                    fe = updated_cache.file_entry(key)
                    if fe:
                        file_token_sum += fe.tokens
                entry = {
                    "parts": list(mod.parts),
                    "part_count": len(mod.parts),
                    "estimated_tokens": file_token_sum,
                    "files": list(mod.files),
                }
            manifest_modules[mid] = entry
            total_tokens += int(entry["estimated_tokens"])

        # Stages 5–6 — transfer, orphans, sleep, manifest
        result = writer.publish(
            modules_to_build,
            all_active_names=active_names,
            manifest_modules=manifest_modules,
            total_estimated_tokens=total_tokens,
            dry_run=False,
        )
        save_cache(project_root, updated_cache)

        rebuilt = ", ".join(m.module_id for m in result.modules) or "(none)"
        _print(
            f"{Fore.GREEN}Done.{Style.RESET_ALL} Rebuilt: {rebuilt}. "
            f"Transferred: {len(result.transferred)} file(s). "
            f"Orphans removed: {len(result.orphans_removed)}."
        )
        if result.manifest_path:
            _print(f"Manifest: {result.manifest_path}")
        return 0
    finally:
        writer.cleanup_temp()
        _active_writer = None


def main() -> None:
    args = run(cleanup)
    raise SystemExit(run_pipeline(force=args.force, dry_run=args.dry_run))


if __name__ == "__main__":
    main()