#!/usr/bin/env python3
"""
Consolidate CrystFlux core files into a single text file for LLM context.

This script collects core Python files from the CrystFlux v1 codebase,
excluding external integrations, and generates a consolidated text file
optimized for LLM consumption.
"""


import argparse

import re

from datetime import datetime

from pathlib import Path

from typing import Any, Dict, List, Tuple



class FileConsolidator:
    """Consolidates CrystFlux core files into text format for LLM context."""

    def __init__(self, root_dir: Path, output_format: str):
        """
        Initialize the consolidator.

        Args:
            root_dir: Root directory containing core files
            output_format: Output format ('single' or 'layered')
        """
        self.root_dir = root_dir
        self.output_format = output_format
        self.repo_root = Path(__file__).resolve().parents[1]
        self.exclude_files = {
            "__init__.py",
        }
        # Adapter layer uses an explicit allowlist of files to include.
        self.allowed_adapter_files = {
            "crystallizer_impl.py",
            "frame_inspection.py",
            "desire_logger_impl.py",
            "missing_logger_impl.py",
            "dream_context.py",
        }
        self.version = "unknown"
        self.version_source = "pyproject.toml (unknown)"
        self.generated_at = ""
        self.sections = ["core", "boundary", "adapters"]

    def resolve_version(self) -> Tuple[str, str]:
        """
        Resolve version string from the installed package.
        """
        import sys
        try:
            sys.path.insert(0, str(self.repo_root / "src"))
            import crystflux
            return str(crystflux.__version__), "crystflux.__version__"
        except Exception:
            return "unknown", "version-resolution-failed"

    def format_root_display(self) -> str:
        """
        Format root directory for display.

        Returns:
            Relative path from repo root when possible, otherwise absolute path
        """
        try:
            return str(self.root_dir.relative_to(self.repo_root))
        except ValueError:
            return str(self.root_dir)

    def build_file_list(
        self, files_by_section: Dict[str, List[Path]], section: str | None
    ) -> List[str]:
        """
        Build a section-scoped file list for header metadata.

        Args:
            files_by_section: Files organized by section
            section: Optional section to scope listing, None for all sections

        Returns:
            List of header lines for file listing
        """
        lines = ["# Files:"]
        sections = [section] if section else self.sections

        for name in sections:
            section_files = files_by_section.get(name, [])
            if not section_files:
                continue
            lines.append(f"# [{name}]")
            for file_path in section_files:
                relative_path = file_path.relative_to(self.root_dir)
                lines.append(f"# - {relative_path}")

        return lines

    def build_meta_block(
        self, files_by_section: Dict[str, List[Path]], output_format: str, section: str | None
    ) -> List[str]:
        """
        Build the metadata header block.

        Args:
            files_by_section: Files organized by section
            output_format: Output format label
            section: Optional section label for layered output

        Returns:
            List of header lines for metadata block
        """
        lines = [
            "# === CrystFlux Consolidated Context ===",
            f"# Version: {self.version}",
            f"# Source: {self.version_source}",
            f"# Generated: {self.generated_at}",
            f"# Format: {output_format}",
        ]
        if section:
            lines.append(f"# Layer: {section}")
        lines.append(f"# Root: {self.format_root_display()}")
        lines.extend(self.build_file_list(files_by_section, section))
        return lines

    def collect_files(self) -> Dict[str, List[Path]]:
        """
        Collect core files organized by section.

        Returns:
            Dictionary mapping section names to lists of file paths
        """
        files: Dict[str, List[Path]] = {}

        for section in self.sections:
            section_dir = self.root_dir / section
            section_files: List[Path] = []

            if not section_dir.exists():
                print(f"Warning: Section directory {section_dir} not found")
                continue

            for file_path in sorted(section_dir.glob("*.py")):
                if file_path.name in self.exclude_files:
                    continue
                if section == "adapters" and file_path.name not in self.allowed_adapter_files:
                    continue
                section_files.append(file_path)

            files[section] = section_files

        return files

    def read_file_safely(self, file_path: Path) -> str:
        """
        Safely read file content with error handling.

        Args:
            file_path: Path to the file to read

        Returns:
            File content or error message
        """
        try:
            return file_path.read_text(encoding="utf-8")
        except Exception as e:
            print(f"Warning: Could not read {file_path}: {e}")
            return f"# [File read error: {e}]"

    def generate_single_output(self, files: Dict[str, List[Path]]) -> str:
        """
        Generate single consolidated text file.

        Args:
            files: Dictionary of files by section

        Returns:
            Consolidated text content
        """
        output_lines: List[str] = self.build_meta_block(files, "single", None)
        output_lines.append("")
        total_files = 0

        for section in self.sections:
            if section not in files or not files[section]:
                continue

            output_lines.append(f"## {section.title()}")
            output_lines.append("")

            for file_path in files[section]:
                relative_path = file_path.relative_to(self.root_dir)
                output_lines.append(f"# File: {relative_path}")

                content = self.read_file_safely(file_path)
                output_lines.append(content)
                output_lines.append("")  # Empty line between files
                total_files += 1

        self.file_count = total_files
        return "\n".join(output_lines)

    def generate_layered_output(self, files: Dict[str, List[Path]]) -> Dict[str, str]:
        """
        Generate layered text files by section.

        Args:
            files: Dictionary of files by section

        Returns:
            Dictionary mapping section names to file contents
        """
        outputs: Dict[str, str] = {}
        total_files = 0

        for section in self.sections:
            if section not in files or not files[section]:
                continue

            output_lines: List[str] = self.build_meta_block(files, "layered", section)
            output_lines.append("")
            output_lines.append(f"## {section.title()}")
            output_lines.append("")

            for file_path in files[section]:
                relative_path = file_path.relative_to(self.root_dir)
                output_lines.append(f"# File: {relative_path}")

                content = self.read_file_safely(file_path)
                output_lines.append(content)
                output_lines.append("")  # Empty line between files
                total_files += 1

            outputs[section] = "\n".join(output_lines)

        self.file_count = total_files
        return outputs

    def estimate_tokens(self, text: str) -> Tuple[int, int]:
        """
        Estimate token count range for the given text.

        Args:
            text: Text to analyze

        Returns:
            Tuple of (min_tokens, max_tokens)
        """
        breakdown = self.estimate_token_breakdown(text)
        return breakdown["total_min"], breakdown["total_max"]

    def estimate_token_breakdown(self, text: str) -> Dict[str, int]:
        """
        Estimate token count with a word-level (Latin) and char-level (Unicode) breakdown.

        Args:
            text: Text to analyze

        Returns:
            Dictionary containing breakdown and totals
        """
        # Estimate ASCII alpha tokens via word chunks; others by char ratios.
        alpha_chunks = re.findall(r"[A-Za-z]+", text)
        alpha_chars = sum(len(chunk) for chunk in alpha_chunks)
        alpha_tokens = sum((len(chunk) + 3) // 4 for chunk in alpha_chunks)

        total_chars = len(text)
        whitespace_chars = sum(1 for ch in text if ch.isspace())
        non_alpha_chars = total_chars - alpha_chars
        non_alpha_non_ws = non_alpha_chars - whitespace_chars

        char_min = int((non_alpha_non_ws / 3.0) + (whitespace_chars / 10.0))
        char_max = int((non_alpha_non_ws / 1.5) + (whitespace_chars / 4.0))
        total_min = int(alpha_tokens + char_min)
        total_max = int(alpha_tokens + char_max)

        return {
            "word_tokens": int(alpha_tokens),
            "char_tokens_min": char_min,
            "char_tokens_max": char_max,
            "total_min": total_min,
            "total_max": total_max,
        }

    def calculate_statistics(self, content: str) -> Dict[str, Any]:
        """
        Calculate statistics for the generated content.

        Args:
            content: Generated text content

        Returns:
            Dictionary containing statistics
        """
        byte_count = len(content.encode("utf-8"))
        breakdown = self.estimate_token_breakdown(content)

        return {
            "byte_count": byte_count,
            "min_tokens": breakdown["total_min"],
            "max_tokens": breakdown["total_max"],
            "word_tokens": breakdown["word_tokens"],
            "char_tokens_min": breakdown["char_tokens_min"],
            "char_tokens_max": breakdown["char_tokens_max"],
            "file_count": getattr(self, "file_count", 0),
        }

    def display_statistics(
        self, stats: Dict[str, Any], output_path: str, files_by_section: Dict[str, List[Path]]
    ):
        """
        Display statistics in English.

        Args:
            stats: Statistics dictionary
            output_path: Output file or directory path
            files_by_section: Files organized by section
        """
        print("\n=== Statistics ===")
        print("Note: Token counts are rough estimates.")
        print(f"Output: {output_path}")
        print(f"Byte count: {stats['byte_count']:,} bytes")
        print(f"Files processed: {stats['file_count']} files")
        print("")
        print("TOTAL")
        print(f"  {'Tokens':<18}: {stats['min_tokens']:,}-{stats['max_tokens']:,}")
        print(f"  {'Word (Latin)':<18}: {stats['word_tokens']:,}")
        print(
            f"  {'Char (Unicode)':<18}: {stats['char_tokens_min']:,}-{stats['char_tokens_max']:,}"
        )

        # Show breakdown by section
        for section in self.sections:
            if section in files_by_section:
                count = len(files_by_section[section])
                if count > 0:
                    print("")
                    print(section.title())
                    print(f"  {'Files':<18}: {count}")

    def display_layered_statistics(
        self, outputs: Dict[str, str], files_by_section: Dict[str, List[Path]], output_dir: Path
    ):
        """
        Display statistics for layered output format.

        Args:
            outputs: Dictionary of section contents
            files_by_section: Files organized by section
            output_dir: Output directory path
        """
        total_bytes = sum(len(content.encode("utf-8")) for content in outputs.values())
        total_min_tokens = sum(self.estimate_tokens(content)[0] for content in outputs.values())
        total_max_tokens = sum(self.estimate_tokens(content)[1] for content in outputs.values())
        total_files = sum(len(files_by_section.get(section, [])) for section in files_by_section)

        print("\n=== Statistics ===")
        print("Note: Token counts are rough estimates.")
        print(f"Output directory: {output_dir}")
        print(f"Total byte count: {total_bytes:,} bytes")
        print(f"Files processed: {total_files} files")
        print("")
        print("TOTAL")
        print(f"  {'Tokens':<18}: {total_min_tokens:,}-{total_max_tokens:,}")

        for section in self.sections:
            if section in outputs:
                content = outputs[section]
                byte_count = len(content.encode("utf-8"))
                breakdown = self.estimate_token_breakdown(content)
                file_count = len(files_by_section.get(section, []))

                print("")
                print(section.title())
                print(f"  {'Tokens':<18}: {breakdown['total_min']:,}-{breakdown['total_max']:,}")
                print(f"  {'Word (Latin)':<18}: {breakdown['word_tokens']:,}")
                print(
                    f"  {'Char (Unicode)':<18}: "
                    f"{breakdown['char_tokens_min']:,}-{breakdown['char_tokens_max']:,}"
                )
                print(f"  {'Bytes':<18}: {byte_count:,}")
                print(f"  {'Files':<18}: {file_count}")

    def run(self, output_path: Path) -> bool:
        """
        Run the consolidation process.

        Args:
            output_path: Output file or directory path

        Returns:
            True if successful, False otherwise
        """
        if not self.root_dir.exists():
            print(f"Error: Root directory {self.root_dir} does not exist")
            return False

        print("Collecting files...")
        files = self.collect_files()
        total_files = sum(len(section_files) for section_files in files.values())

        if total_files == 0:
            print("No files found to process")
            return False

        print(f"Found {total_files} files")
        self.version, self.version_source = self.resolve_version()
        print(f"Detected version: {self.version} (source: {self.version_source})")
        self.generated_at = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        if self.output_format == "single":
            print("Generating single file...")
            content = self.generate_single_output(files)

            try:
                output_path.parent.mkdir(parents=True, exist_ok=True)
                output_path.write_text(content, encoding="utf-8")
                print(f"Generated: {output_path}")

                stats = self.calculate_statistics(content)
                self.display_statistics(stats, str(output_path), files)
                return True
            except Exception as e:
                print(f"Error writing output file: {e}")
                return False

        elif self.output_format == "layered":
            print("Generating layered files...")
            outputs = self.generate_layered_output(files)

            try:
                output_path.mkdir(parents=True, exist_ok=True)

                for section, content in outputs.items():
                    section_file = output_path / f"{section}.txt"
                    section_file.write_text(content, encoding="utf-8")
                    print(f"  Generated: {section_file}")

                self.display_layered_statistics(outputs, files, output_path)
                return True
            except Exception as e:
                print(f"Error writing output files: {e}")
                return False

        return False


def main():
    """Main entry point for the consolidation script."""
    parser = argparse.ArgumentParser(description="Consolidate CrystFlux core files for LLM context")
    parser.add_argument(
        "--root",
        type=Path,
        default=Path("../src/crystflux/v1"),
        help="Root directory containing core files (default: ../src/crystflux/v1)",
    )
    parser.add_argument(
        "--format",
        choices=["single", "layered"],
        default="single",
        help="Output format: single file or layered files",
    )
    parser.add_argument("--output", type=Path, help="Output file for single format")
    parser.add_argument("--output-dir", type=Path, help="Output directory for layered format")

    args = parser.parse_args()

    # Validate arguments
    if args.format == "single" and not args.output:
        parser.error("--output required for single format")
    if args.format == "layered" and not args.output_dir:
        parser.error("--output-dir required for layered format")

    # Create consolidator and run
    consolidator = FileConsolidator(args.root, args.format)
    output_path = args.output if args.format == "single" else args.output_dir

    success = consolidator.run(output_path)

    if success:
        print("\nDone!")
    else:
        print("\nFailed!")
        exit(1)


if __name__ == "__main__":
    main()
