import json
import click
from browser_cli.commands import _handle
from rich.console import Console

console = Console()


@click.group("storage")
def storage_group():
    """Read and write the page's localStorage / sessionStorage."""


@storage_group.command("get")
@click.argument("key", required=False)
@click.option("--type", "store_type", type=click.Choice(["local", "session"]), default="local", show_default=True)
@click.option("--tab", "tab_id", type=int, default=None, help="Tab ID (default: active tab)")
def storage_get(key, store_type, tab_id):
    """Get a localStorage KEY (or dump all keys if omitted)."""
    result = _handle("storage.get", {"key": key, "type": store_type, "tabId": tab_id})
    if result is None:
        console.print("[dim]null[/dim]")
    elif isinstance(result, dict):
        console.print(json.dumps(result, indent=2))
    else:
        console.print(str(result))


@storage_group.command("set")
@click.argument("key")
@click.argument("value")
@click.option("--type", "store_type", type=click.Choice(["local", "session"]), default="local", show_default=True)
@click.option("--tab", "tab_id", type=int, default=None, help="Tab ID (default: active tab)")
def storage_set(key, value, store_type, tab_id):
    """Set localStorage KEY to VALUE."""
    _handle("storage.set", {"key": key, "value": value, "type": store_type, "tabId": tab_id})
    console.print(f"[green]Set[/green] {store_type}[{key!r}] = {value!r}")