import typer
import yaml
import json
from pathlib import Path
from typing import Optional
from rich.console import Console

app = typer.Typer(help="Configuration management commands")
console = Console()

@app.command("s3-layout-add")
def s3_layout_add(
    name: str = typer.Option(..., "--name", "-n", help="Name of the layout"),
    layout_json: str = typer.Option(..., "--layout", "-l", help="JSON string of the layout (e.g. '{\"depot\": [\"sub1\"]}')"),
    project_path: Path = typer.Argument(Path("."), help="Project root directory"),
):
    """
    Add an S3 storage layout to storage.yaml.
    """
    conf_dir = project_path / "conf"
    params_path = conf_dir / "discovery_params.yaml"
    storage_path = conf_dir / "storage.yaml"

    if not params_path.exists():
        # Fallback to legacy
        legacy_dir = project_path / "generators" / "datacubes"
        if (legacy_dir / "discovery_params.yaml").exists():
            params_path = legacy_dir / "discovery_params.yaml"
            storage_path = legacy_dir / "storage.yaml"
        else:
            console.print(f"[red]Error: Project configuration not found at {params_path}.[/red]")
            raise typer.Exit(1)

    try:
        new_layout = json.loads(layout_json)
    except json.JSONDecodeError as e:
        console.print(f"[red]Error: Invalid JSON layout: {e}[/red]")
        raise typer.Exit(1)

    try:
        # Load existing YAML or empty dict
        data = {}
        target_file = storage_path if storage_path.exists() else params_path
        
        # If storage.yaml exists, we use it. If not, we use discovery_params.yaml (legacy)
        # However, for new projects we want storage.yaml. 
        # Let's be smart: if 's3_layouts' is in discovery_params.yaml, stick to it.
        # Otherwise, use storage.yaml.
        
        with open(params_path, "r") as f:
            core_data = yaml.safe_load(f) or {}
            
        if "s3_layouts" in core_data:
            target_file = params_path
            data = core_data
        else:
            target_file = storage_path
            if storage_path.exists():
                with open(storage_path, "r") as f:
                    data = yaml.safe_load(f) or {}

        # Update s3_layouts section
        if "s3_layouts" not in data:
            data["s3_layouts"] = {}
        
        data["s3_layouts"][name] = new_layout

        # Write back
        with open(target_file, "w") as f:
            yaml.dump(data, f, sort_keys=False, default_flow_style=False)

        console.print(f"[green]Successfully added S3 layout '{name}' to {target_file.name}[/green]")
    except Exception as e:
        console.print(f"[red]Error updating configuration: {e}[/red]")
        raise typer.Exit(1)
