client/artdag.py

#!/usr/bin/env python3
"""
Art DAG Client

CLI for interacting with the Art DAG L1 server.
"""

import json
import os
import sys
import time
from pathlib import Path

import click
import requests

DEFAULT_SERVER = os.environ.get("ARTDAG_SERVER", "http://localhost:8100")
DEFAULT_L2_SERVER = os.environ.get("ARTDAG_L2", "http://localhost:8200")
CONFIG_DIR = Path.home() / ".artdag"
TOKEN_FILE = CONFIG_DIR / "token.json"


def get_server():
    """Get server URL from env or default."""
    return DEFAULT_SERVER


def get_l2_server():
    """Get L2 server URL."""
    return DEFAULT_L2_SERVER


def load_token() -> dict:
    """Load saved token from config."""
    if TOKEN_FILE.exists():
        with open(TOKEN_FILE) as f:
            return json.load(f)
    return {}


def save_token(token_data: dict):
    """Save token to config."""
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    with open(TOKEN_FILE, "w") as f:
        json.dump(token_data, f, indent=2)
    TOKEN_FILE.chmod(0o600)


def clear_token():
    """Clear saved token."""
    if TOKEN_FILE.exists():
        TOKEN_FILE.unlink()


def get_auth_header() -> dict:
    """Get Authorization header if token exists."""
    token_data = load_token()
    token = token_data.get("access_token")
    if token:
        return {"Authorization": f"Bearer {token}"}
    return {}


def api_get(path: str, auth: bool = False):
    """GET request to server."""
    headers = get_auth_header() if auth else {}
    resp = requests.get(f"{get_server()}{path}", headers=headers)
    resp.raise_for_status()
    return resp.json()


def api_post(path: str, data: dict = None, params: dict = None, auth: bool = False):
    """POST request to server."""
    headers = get_auth_header() if auth else {}
    resp = requests.post(f"{get_server()}{path}", json=data, params=params, headers=headers)
    resp.raise_for_status()
    return resp.json()


@click.group()
@click.option("--server", "-s", envvar="ARTDAG_SERVER", default=DEFAULT_SERVER,
              help="L1 server URL")
@click.option("--l2", envvar="ARTDAG_L2", default=DEFAULT_L2_SERVER,
              help="L2 server URL")
@click.pass_context
def cli(ctx, server, l2):
    """Art DAG Client - interact with L1 rendering server."""
    ctx.ensure_object(dict)
    ctx.obj["server"] = server
    ctx.obj["l2"] = l2
    global DEFAULT_SERVER, DEFAULT_L2_SERVER
    DEFAULT_SERVER = server
    DEFAULT_L2_SERVER = l2


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

@cli.command()
@click.argument("username")
@click.option("--password", "-p", prompt=True, hide_input=True)
def login(username, password):
    """Login to get access token."""
    try:
        resp = requests.post(
            f"{get_l2_server()}/auth/login",
            json={"username": username, "password": password}
        )
        if resp.status_code == 200:
            token_data = resp.json()
            save_token(token_data)
            click.echo(f"Logged in as {token_data['username']}")
            click.echo(f"Token expires: {token_data['expires_at']}")
        else:
            click.echo(f"Login failed: {resp.json().get('detail', 'Unknown error')}", err=True)
            sys.exit(1)
    except requests.RequestException as e:
        click.echo(f"Login failed: {e}", err=True)
        sys.exit(1)


@cli.command()
@click.argument("username")
@click.option("--password", "-p", prompt=True, hide_input=True, confirmation_prompt=True)
@click.option("--email", "-e", default=None, help="Email (optional)")
def register(username, password, email):
    """Register a new account."""
    try:
        resp = requests.post(
            f"{get_l2_server()}/auth/register",
            json={"username": username, "password": password, "email": email}
        )
        if resp.status_code == 200:
            token_data = resp.json()
            save_token(token_data)
            click.echo(f"Registered and logged in as {token_data['username']}")
        else:
            click.echo(f"Registration failed: {resp.json().get('detail', 'Unknown error')}", err=True)
            sys.exit(1)
    except requests.RequestException as e:
        click.echo(f"Registration failed: {e}", err=True)
        sys.exit(1)


@cli.command()
def logout():
    """Logout (clear saved token)."""
    clear_token()
    click.echo("Logged out")


@cli.command()
def whoami():
    """Show current logged-in user."""
    token_data = load_token()
    if not token_data.get("access_token"):
        click.echo("Not logged in")
        return

    try:
        resp = requests.get(
            f"{get_l2_server()}/auth/me",
            headers={"Authorization": f"Bearer {token_data['access_token']}"}
        )
        if resp.status_code == 200:
            user = resp.json()
            click.echo(f"Username: {user['username']}")
            click.echo(f"Created: {user['created_at']}")
            if user.get('email'):
                click.echo(f"Email: {user['email']}")
        else:
            click.echo("Token invalid or expired. Please login again.", err=True)
            clear_token()
    except requests.RequestException as e:
        click.echo(f"Error: {e}", err=True)


# ============ Server Commands ============

@cli.command()
def info():
    """Show server info."""
    data = api_get("/")
    click.echo(f"Server: {get_server()}")
    click.echo(f"Name: {data['name']}")
    click.echo(f"Version: {data['version']}")
    click.echo(f"Cache: {data['cache_dir']}")
    click.echo(f"Runs: {data['runs_count']}")


@cli.command()
@click.argument("recipe")
@click.argument("input_hash")
@click.option("--name", "-n", help="Output name")
@click.option("--wait", "-w", is_flag=True, help="Wait for completion")
def run(recipe, input_hash, name, wait):
    """Start a rendering run. Requires login.

    RECIPE: Effect/recipe to apply (e.g., dog, identity)
    INPUT_HASH: Content hash of input asset
    """
    # Check auth
    token_data = load_token()
    if not token_data.get("access_token"):
        click.echo("Not logged in. Please run: artdag login <username>", err=True)
        sys.exit(1)

    # Resolve named assets
    assets = api_get("/assets")
    if input_hash in assets:
        input_hash = assets[input_hash]
        click.echo(f"Resolved input to: {input_hash[:16]}...")

    data = {
        "recipe": recipe,
        "inputs": [input_hash],
    }
    if name:
        data["output_name"] = name

    try:
        result = api_post("/runs", data, auth=True)
    except requests.HTTPError as e:
        if e.response.status_code == 401:
            click.echo("Authentication failed. Please login again.", err=True)
            sys.exit(1)
        raise
    run_id = result["run_id"]

    click.echo(f"Run started: {run_id}")
    click.echo(f"Status: {result['status']}")

    if wait:
        click.echo("Waiting for completion...")
        while True:
            status = api_get(f"/runs/{run_id}")
            if status["status"] in ("completed", "failed"):
                break
            time.sleep(1)
            click.echo(".", nl=False)
        click.echo()

        if status["status"] == "completed":
            click.echo(f"Completed!")
            click.echo(f"Output: {status['output_hash']}")
        else:
            click.echo(f"Failed: {status.get('error', 'Unknown error')}")


@cli.command("runs")
@click.option("--limit", "-l", default=10, help="Max runs to show")
def list_runs(limit):
    """List all runs."""
    runs = api_get("/runs")

    if not runs:
        click.echo("No runs found.")
        return

    click.echo(f"{'ID':<36}  {'Status':<10}  {'Recipe':<10}  {'Output Hash':<20}")
    click.echo("-" * 80)

    for run in runs[:limit]:
        output = run.get("output_hash", "")[:16] + "..." if run.get("output_hash") else "-"
        click.echo(f"{run['run_id']}  {run['status']:<10}  {run['recipe']:<10}  {output}")


@cli.command()
@click.argument("run_id")
def status(run_id):
    """Get status of a run."""
    try:
        run = api_get(f"/runs/{run_id}")
    except requests.HTTPError as e:
        if e.response.status_code == 404:
            click.echo(f"Run not found: {run_id}")
            return
        raise

    click.echo(f"Run ID: {run['run_id']}")
    click.echo(f"Status: {run['status']}")
    click.echo(f"Recipe: {run['recipe']}")
    click.echo(f"Inputs: {', '.join(run['inputs'])}")
    click.echo(f"Output Name: {run['output_name']}")
    click.echo(f"Created: {run['created_at']}")

    if run.get("completed_at"):
        click.echo(f"Completed: {run['completed_at']}")

    if run.get("output_hash"):
        click.echo(f"Output Hash: {run['output_hash']}")

    if run.get("error"):
        click.echo(f"Error: {run['error']}")


@cli.command()
@click.option("--limit", "-l", default=20, help="Max items to show")
def cache(limit):
    """List cached content."""
    items = api_get("/cache")

    if not items:
        click.echo("Cache is empty.")
        return

    click.echo(f"Cached content ({len(items)} items):")
    for item in items[:limit]:
        click.echo(f"  {item}")


@cli.command()
@click.argument("content_hash")
@click.option("--output", "-o", type=click.Path(), help="Save to file (use - for stdout)")
def view(content_hash, output):
    """View or download cached content.

    Use -o - to pipe to stdout, e.g.: artdag view <hash> -o - | mpv -
    """
    url = f"{get_server()}/cache/{content_hash}"

    try:
        if output == "-":
            # Stream to stdout for piping
            resp = requests.get(url, stream=True)
            resp.raise_for_status()
            for chunk in resp.iter_content(chunk_size=8192):
                sys.stdout.buffer.write(chunk)
        elif output:
            # Download to file
            resp = requests.get(url, stream=True)
            resp.raise_for_status()
            with open(output, "wb") as f:
                for chunk in resp.iter_content(chunk_size=8192):
                    f.write(chunk)
            click.echo(f"Saved to: {output}", err=True)
        else:
            # Get info via GET with stream to check size
            resp = requests.get(url, stream=True)
            resp.raise_for_status()
            size = resp.headers.get("content-length", "unknown")
            content_type = resp.headers.get("content-type", "unknown")
            click.echo(f"Hash: {content_hash}")
            click.echo(f"Size: {size} bytes")
            click.echo(f"Type: {content_type}")
            click.echo(f"URL: {url}")
            resp.close()
    except requests.HTTPError as e:
        if e.response.status_code == 404:
            click.echo(f"Not found: {content_hash}", err=True)
        else:
            raise


@cli.command("import")
@click.argument("filepath", type=click.Path(exists=True))
def import_file(filepath):
    """Import a local file to cache (local server only)."""
    path = str(Path(filepath).resolve())
    result = api_post("/cache/import", params={"path": path})
    click.echo(f"Imported: {result['content_hash']}")


@cli.command()
@click.argument("filepath", type=click.Path(exists=True))
def upload(filepath):
    """Upload a file to cache (works with remote servers). Requires login."""
    # Check auth
    token_data = load_token()
    if not token_data.get("access_token"):
        click.echo("Not logged in. Please run: artdag login <username>", err=True)
        sys.exit(1)

    try:
        with open(filepath, "rb") as f:
            files = {"file": (Path(filepath).name, f)}
            headers = {"Authorization": f"Bearer {token_data['access_token']}"}
            resp = requests.post(f"{get_server()}/cache/upload", files=files, headers=headers)
            if resp.status_code == 401:
                click.echo("Authentication failed. Please login again.", err=True)
                sys.exit(1)
            resp.raise_for_status()
            result = resp.json()
        click.echo(f"Uploaded: {result['content_hash']}")
        click.echo(f"Size: {result['size']} bytes")
    except requests.RequestException as e:
        click.echo(f"Upload failed: {e}", err=True)
        sys.exit(1)


@cli.command()
def assets():
    """List known assets."""
    data = api_get("/assets")
    click.echo("Known assets:")
    for name, hash in data.items():
        click.echo(f"  {name}: {hash[:16]}...")


@cli.command()
@click.argument("run_id")
@click.argument("output_name")
def publish(run_id, output_name):
    """Publish an L1 run to L2 (register ownership). Requires login.

    RUN_ID: The L1 run ID to publish
    OUTPUT_NAME: Name for the registered asset
    """
    # Check auth
    token_data = load_token()
    if not token_data.get("access_token"):
        click.echo("Not logged in. Please run: artdag login <username>", err=True)
        sys.exit(1)

    # Post to L2 server with auth
    try:
        resp = requests.post(
            f"{get_l2_server()}/registry/record-run",
            json={"run_id": run_id, "output_name": output_name},
            headers={"Authorization": f"Bearer {token_data['access_token']}"}
        )
        if resp.status_code == 401:
            click.echo("Authentication failed. Please login again.", err=True)
            sys.exit(1)
        resp.raise_for_status()
    except requests.RequestException as e:
        click.echo(f"Failed to publish: {e}", err=True)
        sys.exit(1)

    result = resp.json()
    click.echo(f"Published to L2!")
    click.echo(f"Asset: {result['asset']['name']}")
    click.echo(f"Hash: {result['asset']['content_hash']}")
    click.echo(f"Activity: {result['activity']['activity_id']}")


if __name__ == "__main__":
    cli()