from __future__ import annotations

import json
import os
import sys
import time
from dataclasses import asdict
from pathlib import Path
from typing import Optional

import click

from .client import Client
from .exceptions import LiteFoldError

_ENV_CHOICES = click.Choice(["accelerated_cpu", "accelerated_gpu"], case_sensitive=False)


def _get_client(ctx: click.Context) -> Client:
    if "client" not in ctx.obj:
        api_key = ctx.obj.get("api_key")
        if not api_key:
            raise click.UsageError("Missing --api-key or LITEFOLD_API_KEY env var.")
        ctx.obj["client"] = Client(api_key=api_key, base_url=ctx.obj.get("base_url", "https://agentsapi.litefold.ai"))
    return ctx.obj["client"]


def _echo_json(obj):
    click.echo(json.dumps(obj, indent=2, default=str))


def _echo_table(rows: list[dict], columns: Optional[list[str]] = None):
    if not rows:
        click.echo("(no results)")
        return
    columns = columns or list(rows[0].keys())
    widths = {c: max(len(c), *(len(str(r.get(c, ""))) for r in rows)) for c in columns}
    header = "  ".join(c.ljust(widths[c]) for c in columns)
    click.echo(header)
    click.echo("  ".join("-" * widths[c] for c in columns))
    for r in rows:
        click.echo("  ".join(str(r.get(c, "")).ljust(widths[c]) for c in columns))


def _output(ctx: click.Context, data, columns: Optional[list[str]] = None):
    if ctx.obj.get("json_output"):
        _echo_json(data if isinstance(data, (list, dict)) else asdict(data) if hasattr(data, "__dataclass_fields__") else data)
    else:
        if isinstance(data, list):
            rows = [asdict(d) if hasattr(d, "__dataclass_fields__") else d for d in data]
            _echo_table(rows, columns)
        elif hasattr(data, "__dataclass_fields__"):
            for k, v in asdict(data).items():
                click.echo(f"{k}: {v}")
        else:
            click.echo(data)


# ── Root ─────────────────────────────────────────────────────────

@click.group()
@click.option("--api-key", envvar="LITEFOLD_API_KEY", default=None, help="API key (or set LITEFOLD_API_KEY).")
@click.option("--base-url", envvar="LITEFOLD_BASE_URL", default="https://agentsapi.litefold.ai", help="API base URL.")
@click.option("--json", "json_output", is_flag=True, default=False, help="Output as JSON.")
@click.pass_context
def cli(ctx, api_key: Optional[str], base_url: str, json_output: bool):
    ctx.ensure_object(dict)
    ctx.obj["json_output"] = json_output
    ctx.obj["api_key"] = api_key
    ctx.obj["base_url"] = base_url


# ── Projects ─────────────────────────────────────────────────────

@cli.group()
def projects():
    """Manage projects."""


@projects.command("list")
@click.pass_context
def projects_list(ctx):
    """List all projects."""
    result = _get_client(ctx).projects.list()
    _output(ctx, result.projects, ["job_name", "description", "created_at"])


@projects.command("get")
@click.argument("name")
@click.pass_context
def projects_get(ctx, name: str):
    """Get project details."""
    result = _get_client(ctx).projects.get(name)
    _output(ctx, result)


@projects.command("create")
@click.argument("name")
@click.option("--description", "-d", default=None, help="Project description.")
@click.pass_context
def projects_create(ctx, name: str, description: Optional[str]):
    """Create a new project."""
    result = _get_client(ctx).projects.create(name, description)
    _output(ctx, result)


@projects.command("delete")
@click.argument("name")
@click.pass_context
def projects_delete(ctx, name: str):
    """Delete a project."""
    success = _get_client(ctx).projects.delete(name)
    if ctx.obj.get("json_output"):
        _echo_json({"success": success})
    else:
        click.echo("Deleted." if success else "Delete failed.")


# ── Files ────────────────────────────────────────────────────────

@cli.group()
def files():
    """Manage files (upload, list, download, delete)."""


@files.command("upload")
@click.argument("paths", nargs=-1, required=True, type=click.Path(exists=True))
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--folder", "-f", default="upload", help="Destination folder.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_upload(ctx, paths: tuple[str, ...], project: str, folder: str, environment: Optional[str]):
    """Upload one or more files."""
    result = _get_client(ctx).files.upload(project, [Path(p) for p in paths], folder, environment=environment)
    _output(ctx, result)


@files.command("list")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--folder", "-f", default="upload", help="Folder to list.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_list(ctx, project: str, folder: str, environment: Optional[str]):
    """List files in a folder."""
    result = _get_client(ctx).files.list(project, folder, environment=environment)
    _output(ctx, result, ["file_name", "folder", "format_type", "created_at"])


@files.command("list-all")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_list_all(ctx, project: str, environment: Optional[str]):
    """List all files across all folders."""
    result = _get_client(ctx).files.list_all(project, environment=environment)
    _output(ctx, result, ["file_name", "folder", "format_type", "created_at"])


@files.command("list-by-format")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--format", "format_type", required=True, help="Format type (e.g. pdb, sdf, csv).")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_list_by_format(ctx, project: str, format_type: str, environment: Optional[str]):
    """List files filtered by format."""
    result = _get_client(ctx).files.list_by_format(project, format_type, environment=environment)
    _output(ctx, result, ["file_name", "folder", "format_type", "created_at"])


@files.command("list-children")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--parent", required=True, help="Parent file name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_list_children(ctx, project: str, parent: str, environment: Optional[str]):
    """List child files of a parent file."""
    result = _get_client(ctx).files.list_children(project, parent, environment=environment)
    _output(ctx, result, ["file_name", "folder", "format_type"])


@files.command("get-content")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--file-name", required=True, help="File name.")
@click.option("--folder", "-f", default="upload", help="Folder.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_get_content(ctx, project: str, file_name: str, folder: str, environment: Optional[str]):
    """Get text content of a file."""
    content = _get_client(ctx).files.get_content(project, file_name, folder, environment=environment)
    if ctx.obj.get("json_output"):
        _echo_json({"content": content})
    else:
        click.echo(content)


@files.command("get-url")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--file-name", required=True, help="File name.")
@click.option("--folder", "-f", default="upload", help="Folder.")
@click.option("--expires-in", default=3600, type=int, help="URL expiry in seconds.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_get_url(ctx, project: str, file_name: str, folder: str, expires_in: int, environment: Optional[str]):
    """Get a signed download URL."""
    result = _get_client(ctx).files.get_url(project, file_name, folder, expires_in, environment=environment)
    _output(ctx, result)


@files.command("download")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--file-name", required=True, help="File name.")
@click.option("--dest", "-o", default=".", type=click.Path(), help="Destination path (file or directory).")
@click.option("--folder", "-f", default="upload", help="Folder.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_download(ctx, project: str, file_name: str, dest: str, folder: str, environment: Optional[str]):
    """Download a file to local disk."""
    saved = _get_client(ctx).files.download_to(project, file_name, Path(dest), folder, environment=environment)
    if ctx.obj.get("json_output"):
        _echo_json({"saved_to": str(saved)})
    else:
        click.echo(f"Saved to {saved}")


@files.command("get-metadata")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--file-name", required=True, help="File name.")
@click.option("--folder", "-f", default="upload", help="Folder.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_get_metadata(ctx, project: str, file_name: str, folder: str, environment: Optional[str]):
    """Get file metadata."""
    result = _get_client(ctx).files.get_metadata(project, file_name, folder, environment=environment)
    _output(ctx, result)


@files.command("update-metadata")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--file-name", required=True, help="File name.")
@click.option("--metadata", "-m", required=True, help="Metadata as JSON string.")
@click.option("--folder", "-f", default="upload", help="Folder.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_update_metadata(ctx, project: str, file_name: str, metadata: str, folder: str, environment: Optional[str]):
    """Update file metadata."""
    meta = json.loads(metadata)
    result = _get_client(ctx).files.update_metadata(project, file_name, meta, folder, environment=environment)
    if ctx.obj.get("json_output"):
        _echo_json(result)
    else:
        click.echo("Updated.")


@files.command("delete")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--file-name", required=True, help="File name.")
@click.option("--folder", "-f", default="upload", help="Folder.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_delete(ctx, project: str, file_name: str, folder: str, environment: Optional[str]):
    """Delete a single file."""
    success = _get_client(ctx).files.delete(project, file_name, folder, environment=environment)
    if ctx.obj.get("json_output"):
        _echo_json({"success": success})
    else:
        click.echo("Deleted." if success else "Delete failed.")


@files.command("delete-batch")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--file-names", required=True, help="Comma-separated file names.")
@click.option("--folder", "-f", default="upload", help="Folder.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_delete_batch(ctx, project: str, file_names: str, folder: str, environment: Optional[str]):
    """Delete multiple files at once."""
    names = [n.strip() for n in file_names.split(",")]
    result = _get_client(ctx).files.delete_batch(project, names, folder, environment=environment)
    if ctx.obj.get("json_output"):
        _echo_json(result)
    else:
        click.echo(f"Deleted {len(names)} file(s).")


@files.command("upsync")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--folder", "-f", required=True, help="Folder to sync.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_upsync(ctx, project: str, folder: str, environment: Optional[str]):
    """Sync files up to cloud storage."""
    result = _get_client(ctx).files.upsync(project, folder, environment=environment)
    if ctx.obj.get("json_output"):
        _echo_json(result)
    else:
        click.echo("Upsync complete.")


@files.command("downsync")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--folder", "-f", required=True, help="Folder to sync.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def files_downsync(ctx, project: str, folder: str, environment: Optional[str]):
    """Sync files down from cloud storage."""
    result = _get_client(ctx).files.downsync(project, folder, environment=environment)
    if ctx.obj.get("json_output"):
        _echo_json(result)
    else:
        click.echo("Downsync complete.")


# ── Lab ──────────────────────────────────────────────────────────

@cli.group()
def lab():
    """LiteFold Lab — structure prediction, docking, SBDD, and protein design."""


# ── Lab > Structure ──────────────────────────────────────────────

@lab.group()
def structure():
    """Structure prediction jobs."""


@structure.command("submit")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name for this prediction.")
@click.option("--file-names", required=True, help="Comma-separated FASTA/YAML file names.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def structure_submit(ctx, project: str, job_name: str, file_names: str, environment: Optional[str]):
    """Submit a structure prediction job."""
    names = [n.strip() for n in file_names.split(",")]
    result = _get_client(ctx).lab.structure.submit(project, job_name, names, environment=environment)
    _output(ctx, result)


@structure.command("list-jobs")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def structure_list_jobs(ctx, project: str, environment: Optional[str]):
    """List all structure prediction jobs."""
    result = _get_client(ctx).lab.structure.list_jobs(project, environment=environment)
    _output(ctx, result, ["job_id", "job_name", "status", "total_files", "completed", "failed"])


@structure.command("status")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def structure_status(ctx, project: str, job_name: str, environment: Optional[str]):
    """Get status of a structure prediction job."""
    result = _get_client(ctx).lab.structure.get_status(project, job_name, environment=environment)
    _output(ctx, result)


@structure.command("result")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--file-name", required=True, help="File name to get result for.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def structure_result(ctx, project: str, job_name: str, file_name: str, environment: Optional[str]):
    """Get the prediction result for a specific file."""
    result = _get_client(ctx).lab.structure.get_result(project, job_name, file_name, environment=environment)
    _output(ctx, result)


@structure.command("poll")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--interval", default=10, type=int, help="Poll interval in seconds.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def structure_poll(ctx, project: str, job_name: str, interval: int, environment: Optional[str]):
    """Poll a structure prediction job until completion."""
    client = _get_client(ctx)
    while True:
        status = client.lab.structure.get_status(project, job_name, environment=environment)
        job_status = status.job_status or "unknown"
        click.echo(f"  status={job_status}  completed={status.completed}/{status.total_files}  failed={status.failed}")
        if job_status in ("completed", "failed"):
            break
        time.sleep(interval)
    if ctx.obj.get("json_output"):
        _echo_json(asdict(status))
    else:
        click.echo(f"\nJob finished with status: {job_status}")


# ── Lab > Docking ────────────────────────────────────────────────

@lab.group()
def docking():
    """Molecular docking jobs."""


@docking.command("find-pockets")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def docking_find_pockets(ctx, project: str, protein: str, environment: Optional[str]):
    """Detect binding pockets on a protein (one-shot)."""
    result = _get_client(ctx).lab.docking.find_pockets(project, protein, environment=environment)
    if ctx.obj.get("json_output"):
        _echo_json(asdict(result))
    else:
        click.echo(f"Found {result.num_pockets} pocket(s):")
        for p in result.pockets:
            click.echo(f"  #{p.pocket_id}  center={p.center}  score={p.score:.3f}  druggability={p.druggability_score}")


@docking.command("submit")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--ligands", required=True, help="Comma-separated ligand file names.")
@click.option("--pocket-coords", required=True, help="Pocket center x,y,z.")
@click.option("--box-sizes", required=True, help="Box dimensions x,y,z.")
@click.option("--num-poses", default=10, type=int, help="Number of docking poses.")
@click.option("--exhaustiveness", default=128, type=int, help="Search exhaustiveness.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def docking_submit(ctx, project: str, job_name: str, protein: str, ligands: str, pocket_coords: str, box_sizes: str, num_poses: int, exhaustiveness: int, environment: Optional[str]):
    """Submit a docking job."""
    ligand_names = [n.strip() for n in ligands.split(",")]
    coords = [float(x.strip()) for x in pocket_coords.split(",")]
    sizes = [float(x.strip()) for x in box_sizes.split(",")]
    result = _get_client(ctx).lab.docking.submit(
        project, job_name, protein, ligand_names, coords, sizes,
        num_poses=num_poses, exhaustiveness=exhaustiveness,
        environment=environment,
    )
    _output(ctx, result)


@docking.command("list-jobs")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", default=None, help="Filter by job name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def docking_list_jobs(ctx, project: str, job_name: Optional[str], environment: Optional[str]):
    """List docking jobs."""
    result = _get_client(ctx).lab.docking.list_jobs(project, job_name, environment=environment)
    _output(ctx, result, ["job_id", "job_name", "status", "total_ligands", "completed", "failed"])


@docking.command("status")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def docking_status(ctx, project: str, job_name: str, protein: str, environment: Optional[str]):
    """Get status of a docking job."""
    result = _get_client(ctx).lab.docking.get_status(project, job_name, protein, environment=environment)
    _output(ctx, result)


@docking.command("result")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--ligand", required=True, help="Ligand name.")
@click.option("--pose-id", default=None, type=int, help="Specific pose ID.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def docking_result(ctx, project: str, job_name: str, protein: str, ligand: str, pose_id: Optional[int], environment: Optional[str]):
    """Get docking result for a specific ligand."""
    result = _get_client(ctx).lab.docking.get_result(project, job_name, protein, ligand, pose_id, environment=environment)
    _output(ctx, result)


@docking.command("summary")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def docking_summary(ctx, project: str, job_name: str, protein: str, environment: Optional[str]):
    """Get docking summary for all ligands."""
    result = _get_client(ctx).lab.docking.get_summary(project, job_name, protein, environment=environment)
    _output(ctx, result)


@docking.command("poll")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--interval", default=10, type=int, help="Poll interval in seconds.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def docking_poll(ctx, project: str, job_name: str, protein: str, interval: int, environment: Optional[str]):
    """Poll a docking job until completion."""
    client = _get_client(ctx)
    while True:
        status = client.lab.docking.get_status(project, job_name, protein, environment=environment)
        job_status = status.job_status or "unknown"
        click.echo(f"  status={job_status}  completed={status.completed}/{status.total_ligands}  failed={status.failed}")
        if job_status in ("completed", "failed"):
            break
        time.sleep(interval)
    if ctx.obj.get("json_output"):
        _echo_json(asdict(status))
    else:
        click.echo(f"\nJob finished with status: {job_status}")


# ── Lab > SBDD ───────────────────────────────────────────────────

@lab.group()
def sbdd():
    """Structure-based drug design jobs."""


@sbdd.command("submit")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--pocket-coords", required=True, help="Pocket center x,y,z.")
@click.option("--num-mols", default=100, type=int, help="Number of molecules to generate.")
@click.option("--radius", default=13, type=int, help="Pocket radius.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def sbdd_submit(ctx, project: str, job_name: str, protein: str, pocket_coords: str, num_mols: int, radius: int, environment: Optional[str]):
    """Submit an SBDD job."""
    coords = [float(x.strip()) for x in pocket_coords.split(",")]
    result = _get_client(ctx).lab.sbdd.submit(project, job_name, protein, coords, num_mols, radius, environment=environment)
    _output(ctx, result)


@sbdd.command("list-jobs")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", default=None, help="Filter by job name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def sbdd_list_jobs(ctx, project: str, job_name: Optional[str], environment: Optional[str]):
    """List SBDD jobs."""
    result = _get_client(ctx).lab.sbdd.list_jobs(project, job_name, environment=environment)
    _output(ctx, result, ["job_id", "job_name", "status", "total_molecules", "completed", "failed"])


@sbdd.command("status")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def sbdd_status(ctx, project: str, job_name: str, protein: str, environment: Optional[str]):
    """Get status of an SBDD job."""
    result = _get_client(ctx).lab.sbdd.get_status(project, job_name, protein, environment=environment)
    _output(ctx, result)


@sbdd.command("results")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--limit", default=50, type=int, help="Max molecules to return.")
@click.option("--offset", default=0, type=int, help="Offset for pagination.")
@click.option("--sort-by", default="vina_score", help="Sort field.")
@click.option("--sort-order", default="asc", help="Sort order (asc/desc).")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def sbdd_results(ctx, project: str, job_name: str, protein: str, limit: int, offset: int, sort_by: str, sort_order: str, environment: Optional[str]):
    """List generated molecules from an SBDD job."""
    result = _get_client(ctx).lab.sbdd.get_results(project, job_name, protein, limit=limit, offset=offset, sort_by=sort_by, sort_order=sort_order, environment=environment)
    if ctx.obj.get("json_output"):
        _echo_json(asdict(result))
    else:
        click.echo(f"Total: {result.total}  (showing {result.offset}-{result.offset + len(result.molecules)})")
        _output(ctx, result.molecules, ["id", "smiles", "vina_score", "qed", "logp", "validity"])


@sbdd.command("molecule")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--mol-id", required=True, help="Molecule ID.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def sbdd_molecule(ctx, project: str, job_name: str, mol_id: str, environment: Optional[str]):
    """Get detail for a specific generated molecule."""
    result = _get_client(ctx).lab.sbdd.get_molecule(project, job_name, mol_id, environment=environment)
    _output(ctx, result)


@sbdd.command("summary")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def sbdd_summary(ctx, project: str, job_name: str, protein: str, environment: Optional[str]):
    """Get summary stats for an SBDD job."""
    result = _get_client(ctx).lab.sbdd.get_summary(project, job_name, protein, environment=environment)
    _output(ctx, result)


@sbdd.command("poll")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--interval", default=10, type=int, help="Poll interval in seconds.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def sbdd_poll(ctx, project: str, job_name: str, protein: str, interval: int, environment: Optional[str]):
    """Poll an SBDD job until completion."""
    client = _get_client(ctx)
    while True:
        status = client.lab.sbdd.get_status(project, job_name, protein, environment=environment)
        job_status = status.job_status or "unknown"
        click.echo(f"  status={job_status}  generated={status.total_generated}/{status.total_requested}  completed={status.completed}  failed={status.failed}")
        if job_status in ("completed", "failed"):
            break
        time.sleep(interval)
    if ctx.obj.get("json_output"):
        _echo_json(asdict(status))
    else:
        click.echo(f"\nJob finished with status: {job_status}")


# ── Lab > Design ─────────────────────────────────────────────────

@lab.group()
def design():
    """Protein design jobs."""


@design.command("validate")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--config", required=True, help="Path to YAML config file.")
@click.option("--ligand", default=None, help="Ligand file name (optional).")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def design_validate(ctx, project: str, protein: str, config: str, ligand: Optional[str], environment: Optional[str]):
    """Validate a design configuration."""
    config_content = Path(config).read_text()
    result = _get_client(ctx).lab.design.validate(project, protein, config_content, ligand, environment=environment)
    _output(ctx, result)


@design.command("submit")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--model-name", required=True, help="Design model name.")
@click.option("--config-name", required=True, help="Design config name.")
@click.option("--protein", required=True, help="Protein file name.")
@click.option("--config", required=True, help="Path to YAML config file.")
@click.option("--ligand", default=None, help="Ligand file name (optional).")
@click.option("--num-designs", default=None, type=int, help="Number of designs.")
@click.option("--validate-first", is_flag=True, default=False, help="Validate config before submitting.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def design_submit(ctx, project: str, job_name: str, model_name: str, config_name: str, protein: str, config: str, ligand: Optional[str], num_designs: Optional[int], validate_first: bool, environment: Optional[str]):
    """Submit a protein design job."""
    config_content = Path(config).read_text()
    client = _get_client(ctx)
    if validate_first:
        result = client.lab.design.validate_and_submit(
            project, job_name, model_name, config_name, protein, config_content,
            ligand_name=ligand, num_designs=num_designs, environment=environment,
        )
    else:
        result = client.lab.design.submit(
            project, job_name, model_name, config_name, protein, config_content,
            ligand_name=ligand, num_designs=num_designs, environment=environment,
        )
    _output(ctx, result)


@design.command("list-jobs")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", default=None, help="Filter by job name.")
@click.option("--status", "status_filter", default=None, help="Filter by status.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def design_list_jobs(ctx, project: str, job_name: Optional[str], status_filter: Optional[str], environment: Optional[str]):
    """List protein design jobs."""
    result = _get_client(ctx).lab.design.list_jobs(project, job_name, status_filter, environment=environment)
    _output(ctx, result, ["id", "job_name", "design_model_name", "design_config_name", "status", "total_designs"])


@design.command("status")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--model-name", required=True, help="Design model name.")
@click.option("--config-name", required=True, help="Design config name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def design_status(ctx, project: str, job_name: str, model_name: str, config_name: str, environment: Optional[str]):
    """Get status of a protein design job."""
    result = _get_client(ctx).lab.design.get_status(project, job_name, model_name, config_name, environment=environment)
    _output(ctx, result)


@design.command("content")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--model-name", required=True, help="Design model name.")
@click.option("--config-name", required=True, help="Design config name.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def design_content(ctx, project: str, job_name: str, model_name: str, config_name: str, environment: Optional[str]):
    """Get design results/content."""
    result = _get_client(ctx).lab.design.get_content(project, job_name, model_name, config_name, environment=environment)
    _output(ctx, result)


@design.command("poll")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--job-name", "-j", required=True, help="Job name.")
@click.option("--model-name", required=True, help="Design model name.")
@click.option("--config-name", required=True, help="Design config name.")
@click.option("--interval", default=10, type=int, help="Poll interval in seconds.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def design_poll(ctx, project: str, job_name: str, model_name: str, config_name: str, interval: int, environment: Optional[str]):
    """Poll a design job until completion."""
    client = _get_client(ctx)
    while True:
        status = client.lab.design.get_status(project, job_name, model_name, config_name, environment=environment)
        job_status = status.status or "unknown"
        click.echo(f"  status={job_status}")
        if job_status in ("completed", "failed"):
            break
        time.sleep(interval)
    if ctx.obj.get("json_output"):
        _echo_json(asdict(status))
    else:
        click.echo(f"\nJob finished with status: {job_status}")


# ── Rosalind ─────────────────────────────────────────────────────

@cli.group()
def rosalind():
    """Interact with Rosalind agent chats."""


@rosalind.command("create-chat")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.option("--description", "-d", default="", help="Chat description.")
@click.option("--environment", "-e", type=_ENV_CHOICES, default=None, help="Environment (accelerated_cpu/accelerated_gpu).")
@click.pass_context
def rosalind_create_chat(ctx, project: str, chat_name: str, description: str, environment: Optional[str]):
    """Create a new chat session."""
    name = _get_client(ctx).rosalind.create_chat(project, chat_name, description, environment=environment)
    if ctx.obj.get("json_output"):
        _echo_json({"chat_name": name})
    else:
        click.echo(f"Chat created: {name}")


@rosalind.command("send-message")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.option("--message", "-m", required=True, help="Message to send.")
@click.option("--mode", default="pro", help="Model mode (pro/max).")
@click.option("--depth", default="balanced", help="Depth (quick/balanced/ultra).")
@click.option("--force-fresh-session", is_flag=True, default=False, help="Force a fresh session.")
@click.pass_context
def rosalind_send_message(ctx, project: str, chat_name: str, message: str, mode: str, depth: str, force_fresh_session: bool):
    """Send a message to a chat (fire-and-forget)."""
    success = _get_client(ctx).rosalind.send_message(
        project, chat_name, message, mode, depth,
        force_fresh_session=force_fresh_session,
    )
    if ctx.obj.get("json_output"):
        _echo_json({"success": success})
    else:
        click.echo("Message sent." if success else "Send failed.")


@rosalind.command("cancel-chat")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.pass_context
def rosalind_cancel_chat(ctx, project: str, chat_name: str):
    """Cancel a running chat task."""
    success = _get_client(ctx).rosalind.cancel_chat(project, chat_name)
    if ctx.obj.get("json_output"):
        _echo_json({"success": success})
    else:
        click.echo("Cancelled." if success else "Nothing to cancel.")


@rosalind.command("list-chats")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--limit", "-l", default=None, type=int, help="Max chats to show.")
@click.pass_context
def rosalind_list_chats(ctx, project: str, limit: Optional[int]):
    """List chats for a project."""
    result = _get_client(ctx).rosalind.list_chats(project, limit)
    _output(ctx, result, ["chat_name", "device", "total_input_tokens", "total_output_tokens", "total_cost_usd"])


@rosalind.command("chat-status")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.pass_context
def rosalind_chat_status(ctx, project: str, chat_name: str):
    """Check if a chat's agent is still running."""
    result = _get_client(ctx).rosalind.chat_status(project, chat_name)
    _output(ctx, result)


@rosalind.command("list-blocks")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.pass_context
def rosalind_list_blocks(ctx, project: str, chat_name: str):
    """Show total number of blocks in a chat."""
    total = _get_client(ctx).rosalind.list_blocks(project, chat_name)
    if ctx.obj.get("json_output"):
        _echo_json({"total": total})
    else:
        click.echo(f"Total blocks: {total}")


@rosalind.command("get-blocks")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.option("--start", default=0, type=int, help="Start index.")
@click.option("--end", default=None, type=int, help="End index.")
@click.pass_context
def rosalind_get_blocks(ctx, project: str, chat_name: str, start: int, end: Optional[int]):
    """Get blocks from chat history."""
    blocks = _get_client(ctx).rosalind.get_blocks(project, chat_name, start, end)
    if ctx.obj.get("json_output"):
        _echo_json([asdict(b) for b in blocks])
    else:
        for b in blocks:
            if b.assistant_message:
                click.secho(f"\n[block {b.index}] assistant:", fg="green", bold=True)
                click.echo(b.assistant_message)
            elif b.content_type == "thinking":
                click.secho(f"\n[block {b.index}] thinking:", fg="yellow", dim=True)
                click.echo(b.text)
            elif b.content_type == "tool_use":
                click.secho(f"[block {b.index}] tool_use: {b.tool_name}", fg="cyan")
            elif b.content_type == "mcp_result":
                label = "output" if b.is_output_panel else "result"
                click.secho(f"[block {b.index}] {label}: {b.tool_name}", fg="blue")
                click.echo(str(b.tool_result)[:500] if b.tool_result else "")
            elif b.content_type == "user_message":
                click.secho(f"\n[block {b.index}] user:", fg="magenta", bold=True)
                click.echo(b.text or "")


@rosalind.command("get-bash-outputs")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.option("--start", default=0, type=int, help="Start index.")
@click.option("--end", default=None, type=int, help="End index.")
@click.pass_context
def rosalind_get_bash_outputs(ctx, project: str, chat_name: str, start: int, end: Optional[int]):
    """Extract bash command inputs/outputs from chat."""
    results = _get_client(ctx).rosalind.get_bash_outputs(project, chat_name, start, end)
    if ctx.obj.get("json_output"):
        _echo_json([asdict(b) for b in results])
    else:
        if not results:
            click.echo("No bash outputs.")
            return
        for b in results:
            click.secho(f"$ {b.command}", fg="green")
            if b.exit_code is not None:
                click.echo(f"  exit_code: {b.exit_code}")
            if b.output:
                click.echo(b.output)
            click.echo()


@rosalind.command("send-and-wait")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.option("--message", "-m", required=True, help="Message to send.")
@click.option("--mode", default="pro", help="Model mode.")
@click.option("--depth", default="balanced", help="Depth.")
@click.option("--poll-interval", default=5, type=int, help="Seconds between status checks.")
@click.pass_context
def rosalind_send_and_wait(ctx, project: str, chat_name: str, message: str, mode: str, depth: str, poll_interval: int):
    """Send a message and wait until the agent finishes, then print blocks."""
    client = _get_client(ctx)
    success = client.rosalind.send_message(project, chat_name, message, mode, depth)
    if not success:
        click.secho("Send failed.", fg="red")
        raise SystemExit(1)

    click.echo("Message sent. Waiting for agent to finish...")
    while True:
        status = client.rosalind.chat_status(project, chat_name)
        if not status.is_running:
            break
        click.echo(f"  running (task={status.task_id})...")
        time.sleep(poll_interval)

    click.echo("Agent finished. Fetching blocks...\n")
    ctx.invoke(rosalind_get_blocks, project=project, chat_name=chat_name, start=0, end=None)


# ── Workspace (sub-group of rosalind) ────────────────────────────

@rosalind.group("workspace")
def rosalind_workspace():
    """Browse workspace files for a chat."""


@rosalind_workspace.command("list-folders")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.pass_context
def workspace_list_folders(ctx, project: str, chat_name: str):
    """List top-level folders in the workspace."""
    ws = _get_client(ctx).rosalind.workspace(project, chat_name)
    folders = ws.list_folders()
    if ctx.obj.get("json_output"):
        _echo_json(folders)
    else:
        for f in folders:
            click.echo(f"  {f}/")
        if not folders:
            click.echo("(no folders)")


@rosalind_workspace.command("list-files")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.option("--folder", "-f", default=None, help="Filter by folder.")
@click.pass_context
def workspace_list_files(ctx, project: str, chat_name: str, folder: Optional[str]):
    """List files in the workspace."""
    ws = _get_client(ctx).rosalind.workspace(project, chat_name)
    result = ws.list_files(folder)
    _output(ctx, result, ["name", "location", "folder", "size"])


@rosalind_workspace.command("get-file-content")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.option("--file-path", required=True, help="File path in workspace.")
@click.pass_context
def workspace_get_file_content(ctx, project: str, chat_name: str, file_path: str):
    """Get the content of a text-viewable workspace file."""
    from .namespaces.workspace import is_text_viewable

    if not is_text_viewable(file_path):
        click.secho(
            f"File '{file_path}' is not a text-viewable format.\n"
            f"Use `litefold rosalind workspace download-file-content` instead.",
            fg="yellow", err=True,
        )
        raise SystemExit(1)

    try:
        ws = _get_client(ctx).rosalind.workspace(project, chat_name)
        fc = ws.get_file_content(file_path)
    except LiteFoldError as e:
        click.secho(f"Error fetching file content: {e}", fg="red", err=True)
        raise SystemExit(1)

    if ctx.obj.get("json_output"):
        _echo_json(asdict(fc))
    else:
        click.echo(fc.content)


@rosalind_workspace.command("download-file-content")
@click.option("--project", "-p", required=True, help="Project name.")
@click.option("--chat-name", "-n", required=True, help="Chat name.")
@click.option("--file-path", required=True, help="File path in workspace.")
@click.option("--save-as", required=True, help="Local path to save the file to.")
@click.pass_context
def workspace_download_file_content(
    ctx, project: str, chat_name: str, file_path: str, save_as: str
):
    """Download a workspace file to a local path."""
    try:
        ws = _get_client(ctx).rosalind.workspace(project, chat_name)
        dest = ws.download_file_content(file_path, save_as)
        click.secho(f"Downloaded to {dest}", fg="green")
    except LiteFoldError as e:
        click.secho(f"Error downloading file: {e}", fg="red", err=True)
        raise SystemExit(1)


# ── Entrypoint ───────────────────────────────────────────────────

def main():
    try:
        cli(standalone_mode=False)
    except LiteFoldError as e:
        click.secho(f"Error: {e}", fg="red", err=True)
        raise SystemExit(1)
    except click.exceptions.Exit:
        pass


if __name__ == "__main__":
    main()
