from __future__ import annotations
import yaml
from pathlib import Path
from typing import List, Dict, Any, Set, Optional
from rich.console import Console
import typer
from pydantic import BaseModel, Field
from collections import defaultdict

from sibi_flux.core import SecureResource
from sibi_flux.utils.security import is_secure_path, is_valid_identifier, is_valid_module_path
from sibi_flux.datacube.generator import generate_datacube_module_code, DatacubeRegistry

console = Console()

class CubeExtensionEntry(BaseModel):
    """Entry for a datacube extension in datacubes.yaml."""
    source: str
    name: str # The class name to use in the application
    module: str # The target module directory in the app
    gen_module: str # The original module path in the global registry
    original_class: Optional[str] = None

class AppDatacubeConfig(BaseModel):
    """Registry of datacube extensions for an application."""
    cubes: List[CubeExtensionEntry] = Field(default_factory=list)

AppDatacubeConfig.model_rebuild()

class CubeExtender(SecureResource):
    """
    Generates app-specific Datacube extensions.
    Hardened with SecureResource for sandboxed I/O.
    """
    
    def __init__(self, **kwargs: Any) -> None:
        super().__init__(**kwargs)

    def create_extensions(self, app_name: str, force: bool = False) -> None:
        """
        Processes the app's datacubes.yaml and generates the extension modules.
        """
        app_dir = self.project_root / app_name
        if not app_dir.exists() or not app_dir.is_dir():
            console.print(f"[red]Error: Application directory '{app_name}' not found.[/red]")
            raise typer.Exit(code=1)
            
        datacubes_dir = app_dir / "datacubes"
        datacubes_dir.mkdir(parents=True, exist_ok=True)
        
        config_path = self.get_secure_path(datacubes_dir / "datacubes.yaml")
        if not config_path.exists():
            console.print(f"[yellow]Info: No Datacube extensions configured for '{app_name}'.[/yellow]")
            return

        # 1. Load Configuration
        with self.open_secure(config_path, "r") as f:
            raw_data = yaml.safe_load(f) or {}
            config = AppDatacubeConfig(**raw_data)

        if not config.cubes:
            console.print(f"[yellow]Info: Empty Datacube list for '{app_name}'.[/yellow]")
            return

        # 2. Group by module name
        cubes_by_module = defaultdict(list)
        for cube in config.cubes:
            cubes_by_module[cube.module].append(cube)

        # 3. Process modules
        found_modules = set()
        for module_name, entries in cubes_by_module.items():
            found_modules.add(module_name)
            self._process_module_group(module_name, entries, datacubes_dir, app_name, force=force)

        # 4. Prune orphans
        if force:
            self._prune_orphaned_modules(datacubes_dir, found_modules)

    def _process_module_group(
        self, 
        module_name: str, 
        entries: List[CubeExtensionEntry], 
        datacubes_dir: Path, 
        app_name: str,
        force: bool = False
    ) -> None:
        """Generates code for a single module package."""
        if not is_valid_identifier(module_name):
             console.print(f"  [yellow]Skipping module '{module_name}': Invalid Python identifier.[/yellow]")
             return
             
        target_dir = datacubes_dir / module_name
        target_dir.mkdir(parents=True, exist_ok=True)
        target_file = self.get_secure_path(target_dir / "__init__.py")

        if target_file.exists() and not force:
            return

        console.print(f"  [blue]Generating module: {module_name} ({len(entries)} cubes)[/blue]")
        
        imports = [
            "from __future__ import annotations",
            "from typing import Optional, List, Dict, Any"
        ]
        
        classes = []
        for entry in entries:
            gen_mod = entry.gen_module
            orig_cls = entry.original_class or entry.name
            app_cls = entry.name
            
            # Security Validation: Ensure valid identifiers for code generation
            if not is_valid_module_path(gen_mod) or not is_valid_identifier(orig_cls) or not is_valid_identifier(app_cls):
                console.print(f"  [yellow]Skipping cube '{app_cls}': Invalid identifier in registry.[/yellow]")
                continue
            
            # Extension boilerplate
            imports.append(f"from {gen_mod} import {orig_cls} as Base{orig_cls}")
            
            class_code = [
                f"class {app_cls}(Base{orig_cls}):",
                "    \"\"\"",
                f"    App-specific extension for {orig_cls}.",
                "    \"\"\"",
                "    pass"
            ]
            classes.append("\n".join(class_code))

        # Build full content
        full_content = imports + ["\n"] + classes
        self.write_text_secure(target_file, "\n".join(full_content))

    def _prune_orphaned_modules(self, datacubes_dir: Path, active_modules: Set[str]) -> None:
        """Removes module directories that are no longer in the registry."""
        for item in datacubes_dir.iterdir():
            if item.is_dir() and item.name not in active_modules and item.name != "__pycache__":
                console.print(f"  [yellow]Pruning orphaned module: {item.name}[/yellow]")
                self._recursive_delete(item)

    def _recursive_delete(self, path: Path) -> None:
        """Securely deletes a directory tree within the sandbox."""
        secure_path = self.get_secure_path(path)
        if secure_path.is_dir():
            for child in secure_path.iterdir():
                self._recursive_delete(child)
            secure_path.rmdir()
        else:
            secure_path.unlink()

def create_cubes(app_name: str, force: bool = False) -> None:
    """Entry point for CLI or legacy calls."""
    extender = CubeExtender()
    extender.create_extensions(app_name, force=force)
