"""Command-line interface for JaxLatt package."""

import typer
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich import box
from rich.text import Text
from rich.align import Align

console = Console()

app = typer.Typer(
    name="jaxlatt",
    help="JAX-based Lattice Field Theory Simulation Package",
    no_args_is_help=True,
)

LOGO = """
     ██╗ █████╗ ██╗  ██╗██╗      █████╗ ████████╗████████╗
     ██║██╔══██╗╚██╗██╔╝██║     ██╔══██╗╚══██╔══╝╚══██╔══╝
     ██║███████║ ╚███╔╝ ██║     ███████║   ██║      ██║   
██   ██║██╔══██║ ██╔██╗ ██║     ██╔══██║   ██║      ██║   
╚█████╔╝██║  ██║██╔╝ ██╗███████╗██║  ██║   ██║      ██║   
 ╚════╝ ╚═╝  ╚═╝╚═╝  ╚═╝╚══════╝╚═╝  ╚═╝   ╚═╝      ╚═╝   
"""


@app.command()
def version():
    """Display the version of JaxLatt."""
    from jaxlatt import __version__

    console.print(
        f"[bold cyan]JaxLatt[/bold cyan] version [bold green]{__version__}[/bold green]"
    )


@app.command()
def info():
    """Display package information and module structure."""
    from jaxlatt import __version__, __author__

    # Print logo (centered)
    logo_text = Text(LOGO, style="bold cyan")
    console.print(Align.center(logo_text))

    # Package info panel (centered)
    info_content = Text.from_markup(
        f"""[bold]Version:[/bold] {__version__}
[bold]Author:[/bold] {__author__}
[bold]Description:[/bold] JAX-based Lattice Field Theory Simulations
[bold]Repository:[/bold] https://github.com/rcalderonb6/jax-latt""",
        justify="center",
    )

    console.print(
        Panel(
            info_content,
            title="[bold cyan]Package Information[/bold cyan]",
            border_style="cyan",
            box=box.ROUNDED,
        )
    )

    # Module structure table (centered)
    table = Table(title="Module Structure", box=box.ROUNDED, border_style="green")
    table.add_column("Module", style="cyan", no_wrap=True)
    table.add_column("Description", style="white")

    table.add_row("jaxlatt.core", "Lattice classes, cosmology, potentials")
    table.add_row(
        "jaxlatt.operators",
        "Field operators (scalar, gauge, coupled, spectral, autodiff)",
    )
    table.add_row(
        "jaxlatt.evolution", "Time integrators (leapfrog, RK4, expanding space)"
    )
    table.add_row(
        "jaxlatt.observables", "Physics observables, spectra, energy tracking"
    )

    console.print(Align.center(table))

    # Quick start example
    example_code = """[dim]from[/dim] jaxlatt.core [dim]import[/dim] create_random_coupled_lattice
[dim]from[/dim] jaxlatt.evolution [dim]import[/dim] coupled_leapfrog_step

[green]# Create lattice[/green]
lattice = create_random_coupled_lattice(key, size=(32,32,32), length=10.0)

[green]# Evolve[/green]
[dim]for[/dim] _ [dim]in[/dim] range(100):
    lattice = coupled_leapfrog_step(lattice, dt=0.01)"""

    console.print(
        Panel(
            example_code,
            title="[bold yellow]Quick Start Example[/bold yellow]",
            border_style="yellow",
            box=box.ROUNDED,
        )
    )


@app.command()
def list_potentials():
    """List available potential functions."""
    console.print("\n[bold cyan]Available Potential Functions[/bold cyan]\n")

    table = Table(box=box.ROUNDED, border_style="magenta")
    table.add_column("Method", style="cyan", no_wrap=True)
    table.add_column("Formula", style="yellow")
    table.add_column("Parameters", style="green")

    table.add_row("ScalarPotential.quadratic(m)", "V(φ) = ½m²φ²", "m: mass")
    table.add_row("ScalarPotential.quartic(m, λ)", "V(φ) = ½m²|φ|² + ¼λ|φ|⁴", "m: mass, λ: coupling")
    table.add_row(
        "ScalarPotential.double_well(μ², λ)",
        "V(φ) = -½μ²φ² + ¼λφ⁴",
        "μ²: quadratic coeff, λ: coupling",
    )
    table.add_row("ScalarPotential.mexican_hat(λ, v)", "V(φ) = λ(|φ|² - v²)²", "λ: coupling, v: VEV")

    console.print(table)

    console.print("\n[bold]Usage:[/bold]")
    console.print(
        "  [dim]from[/dim] jaxlatt.core [dim]import[/dim] ScalarPotential"
    )
    console.print("  V = ScalarPotential.quadratic(m=1.0)")
    console.print("  energy_density = V(field)")
    console.print("  force = V.force(field)  [dim]# F = -dV/dφ via autodiff[/dim]\n")


@app.command()
def check_install():
    """Check that JaxLatt and dependencies are properly installed."""
    console.print("\n[bold cyan]Checking JaxLatt installation...[/bold cyan]\n")

    # Check JAX
    try:
        import jax

        console.print(f"[green]✓[/green] JAX [bold]{jax.__version__}[/bold] installed")
    except ImportError:
        console.print("[red]✗ JAX not found[/red]")
        return

    # Check JaxLatt modules
    modules = ["core", "operators", "evolution", "observables"]
    for module in modules:
        try:
            __import__(f"jaxlatt.{module}")
            console.print(f"[green]✓[/green] jaxlatt.{module} available")
        except ImportError as e:
            console.print(f"[red]✗ jaxlatt.{module} failed: {e}[/red]")

    console.print("\n[bold green]✓ JaxLatt installation OK![/bold green]\n")


def main():
    """Entry point for the CLI."""
    app()


if __name__ == "__main__":
    main()
