"""Anor Cloud CLI"""
import click
from rich.console import Console
from rich.table import Table
from rich.progress import Progress, SpinnerColumn, TextColumn
from typing import Optional, Tuple
import getpass

from . import __version__
from .config import is_logged_in, set_api_key, clear_credentials, get_api_key
from .api import AnorClient, APIError

console = Console()


import functools

def require_auth(func):
    """Decorator to require authentication"""
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if not is_logged_in():
            console.print("[red]Not logged in. Run `anor login` first.[/red]")
            raise SystemExit(1)
        return func(*args, **kwargs)
    return wrapper


@click.group()
@click.version_option(version=__version__, prog_name="anor")
def main():
    """Anor Cloud GPU CLI - Managed GPU compute made simple."""
    pass


# ============ Auth Commands ============

@main.command()
@click.option("--email", "-e", help="Email address")
@click.option("--password", "-p", help="Password")
@click.option("--api-key", "-k", help="Use existing API key")
def login(email: Optional[str], password: Optional[str], api_key: Optional[str]):
    """Login to Anor Cloud"""
    if api_key:
        set_api_key(api_key)
        console.print("[green]API key saved successfully![/green]")
        return

    if not email:
        email = click.prompt("Email")
    if not password:
        password = getpass.getpass("Password: ")

    try:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Logging in...", total=None)
            client = AnorClient()
            result = client.login(email, password)
        
        if "accessToken" in result:
            # Create an API key for CLI use
            client = AnorClient(api_key=result["accessToken"])
            key_result = client.create_api_key("anor-cli")
            set_api_key(key_result["key"])
            console.print(f"[green]Logged in as {email}[/green]")
        else:
            console.print("[red]Login failed: No access token received[/red]")
    except APIError as e:
        console.print(f"[red]Login failed: {e.message}[/red]")
        raise SystemExit(1)


@main.command()
def logout():
    """Logout and clear credentials"""
    clear_credentials()
    console.print("[green]Logged out successfully[/green]")


@main.command()
def whoami():
    """Show current user"""
    if not is_logged_in():
        console.print("[yellow]Not logged in[/yellow]")
        return

    try:
        client = AnorClient()
        user = client.get_me()
        console.print(f"[bold]Email:[/bold] {user.get('email', 'N/A')}")
        console.print(f"[bold]Name:[/bold] {user.get('name', 'N/A')}")
        console.print(f"[bold]ID:[/bold] {user.get('id', 'N/A')}")
    except APIError as e:
        console.print(f"[red]Error: {e.message}[/red]")


# ============ Run Command ============

@main.command()
@click.argument("command", nargs=-1, required=True)
@click.option("--grade", "-g", required=True, type=click.Choice(["standard", "performance", "premium"], case_sensitive=False), help="Compute grade")
@click.option("--gpus", required=True, type=int, help="Number of GPUs")
@click.option("--max-rate", required=True, type=float, help="Max hourly rate ($/GPU/hr)")
@click.option("--name", "-n", help="Job name")
@click.option("--image", "-i", help="Docker image")
@click.option("--workdir", "-w", help="Working directory")
@click.option("--env", "-e", multiple=True, help="Environment variables (KEY=VALUE)")
@require_auth
def run(command: Tuple[str, ...], grade: str, gpus: int, max_rate: float, 
        name: Optional[str], image: Optional[str], workdir: Optional[str], env: Tuple[str, ...]):
    """Submit a compute job"""
    cmd_str = " ".join(command)
    
    env_vars = {}
    for e in env:
        if "=" in e:
            key, value = e.split("=", 1)
            env_vars[key] = value

    try:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Submitting job...", total=None)
            client = AnorClient()
            job = client.create_job(
                command=cmd_str,
                grade=grade,
                gpu_count=gpus,
                max_hourly_rate=max_rate,
                name=name,
                image=image,
                workdir=workdir,
                env_vars=env_vars if env_vars else None,
            )

        console.print("[green]Job submitted successfully![/green]\n")
        console.print(f"[bold]Job ID:[/bold] {job['id']}")
        console.print(f"[bold]Grade:[/bold] {job['grade']}")
        console.print(f"[bold]GPUs:[/bold] {job['gpuCount']}")
        console.print(f"[bold]Max Rate:[/bold] ${job['maxHourlyRate']}/GPU/hr")
        console.print(f"[bold]State:[/bold] {job['state']}")
        console.print()
        console.print(f"[dim]Track with: anor jobs get {job['id']}[/dim]")
        console.print(f"[dim]View logs:  anor jobs logs {job['id']}[/dim]")

    except APIError as e:
        console.print(f"[red]Failed to submit job: {e.message}[/red]")
        raise SystemExit(1)


# ============ Jobs Commands ============

@main.group()
def jobs():
    """Manage compute jobs"""
    pass


@jobs.command("list")
@click.option("--state", "-s", help="Filter by state")
@click.option("--limit", "-l", default=20, help="Number of jobs to show")
@require_auth
def jobs_list(state: Optional[str], limit: int):
    """List your jobs"""
    try:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Fetching jobs...", total=None)
            client = AnorClient()
            result = client.list_jobs(state=state, limit=limit)

        jobs_data = result.get("jobs", [])
        
        if not jobs_data:
            console.print("[yellow]No jobs found[/yellow]")
            return

        table = Table(title="Jobs")
        table.add_column("ID", style="cyan")
        table.add_column("Name")
        table.add_column("Grade")
        table.add_column("GPUs", justify="right")
        table.add_column("State")
        table.add_column("Created")

        for job in jobs_data:
            state_color = {
                "QUEUED": "yellow",
                "ROUTING": "yellow", 
                "STARTING": "blue",
                "RUNNING": "green",
                "SUCCEEDED": "green",
                "FAILED": "red",
                "CANCELLED": "dim",
            }.get(job["state"], "white")
            
            table.add_row(
                job["id"][:8],
                job.get("name") or "-",
                job["grade"],
                str(job["gpuCount"]),
                f"[{state_color}]{job['state']}[/{state_color}]",
                job.get("queuedAt", "")[:19] if job.get("queuedAt") else "-",
            )

        console.print(table)
        console.print(f"\nShowing {len(jobs_data)} of {result.get('total', len(jobs_data))} jobs")

    except APIError as e:
        console.print(f"[red]Error: {e.message}[/red]")
        raise SystemExit(1)


@jobs.command("get")
@click.argument("job_id")
@require_auth
def jobs_get(job_id: str):
    """Get job details"""
    try:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Fetching job...", total=None)
            client = AnorClient()
            job = client.get_job(job_id)

        console.print("\n[bold]Job Details[/bold]")
        console.print("─" * 40)
        console.print(f"[bold]ID:[/bold] {job['id']}")
        console.print(f"[bold]Grade:[/bold] {job['grade']}")
        console.print(f"[bold]GPUs:[/bold] {job['gpuCount']}")
        console.print(f"[bold]Command:[/bold] {job['command']}")
        if job.get("image"):
            console.print(f"[bold]Image:[/bold] {job['image']}")
        console.print(f"[bold]Max Rate:[/bold] ${job['maxHourlyRate']}/GPU/hr")
        if job.get("softLockPrice"):
            console.print(f"[bold]Locked Price:[/bold] ${job['softLockPrice']}/GPU/hr")
        console.print()
        console.print(f"[bold]State:[/bold] {job['state']}")
        if job.get("stateMessage"):
            console.print(f"[bold]Message:[/bold] {job['stateMessage']}")
        console.print()
        if job.get("queuedAt"):
            console.print(f"[bold]Queued:[/bold] {job['queuedAt']}")
        if job.get("startedAt"):
            console.print(f"[bold]Started:[/bold] {job['startedAt']}")
        if job.get("endedAt"):
            console.print(f"[bold]Ended:[/bold] {job['endedAt']}")

    except APIError as e:
        console.print(f"[red]Error: {e.message}[/red]")
        raise SystemExit(1)


@jobs.command("logs")
@click.argument("job_id")
@click.option("--follow", "-f", is_flag=True, help="Follow log output (not yet implemented)")
@require_auth
def jobs_logs(job_id: str, follow: bool):
    """Get job logs"""
    try:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Fetching logs...", total=None)
            client = AnorClient()
            logs = client.get_job_logs(job_id)

        if logs:
            console.print(logs)
        else:
            console.print("[yellow]No logs available yet[/yellow]")

    except APIError as e:
        console.print(f"[red]Error: {e.message}[/red]")
        raise SystemExit(1)


@jobs.command("cancel")
@click.argument("job_id")
@click.confirmation_option(prompt="Are you sure you want to cancel this job?")
@require_auth
def jobs_cancel(job_id: str):
    """Cancel a job"""
    try:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Cancelling job...", total=None)
            client = AnorClient()
            job = client.cancel_job(job_id)

        console.print(f"[green]Job {job_id} cancelled[/green]")

    except APIError as e:
        console.print(f"[red]Error: {e.message}[/red]")
        raise SystemExit(1)


# ============ Pricing Command ============

@main.command()
def pricing():
    """Show current compute pricing"""
    try:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Fetching pricing...", total=None)
            client = AnorClient()
            prices = client.get_pricing()

        table = Table(title="Anor Compute Pricing")
        table.add_column("Grade", style="bold")
        table.add_column("Price ($/GPU/hr)", justify="right")
        table.add_column("Available GPUs", justify="right")

        for tier in prices:
            table.add_row(
                tier["grade"].title(),
                f"${tier['pricePerGpuHour']:.2f}",
                str(tier.get("availableGpus", 0)),
            )

        console.print(table)
        console.print("\n[dim]Prices subject to real-time market conditions.[/dim]")
        console.print("[dim]Use --max-rate to cap your spend.[/dim]")

    except APIError as e:
        console.print(f"[red]Error: {e.message}[/red]")
        raise SystemExit(1)


# ============ Balance Command ============

@main.command()
@require_auth
def balance():
    """Show account balance"""
    try:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Fetching balance...", total=None)
            client = AnorClient()
            data = client.get_balance()

        console.print(f"\n[bold]Available Balance:[/bold] ${data.get('available', 0):.2f}")
        console.print(f"[bold]Pending:[/bold] ${data.get('pending', 0):.2f}")
        console.print(f"[bold]Total:[/bold] ${data.get('total', 0):.2f}")

    except APIError as e:
        console.print(f"[red]Error: {e.message}[/red]")
        raise SystemExit(1)


if __name__ == "__main__":
    main()
