"""Interactive onboarding flow for new users.

Triggered by `ml launch` with no arguments, or callable from `ml setup`.
Creates a task YAML template and optionally writes an AGENTS.md that points
coding agents to the Mithril CLI docs bundled with this package.
"""

from __future__ import annotations

import sys
from importlib.resources import files
from pathlib import Path
from typing import TYPE_CHECKING

from rich.prompt import Confirm

if TYPE_CHECKING:
    from rich.console import Console

TASK_TEMPLATE = """\
# task.yaml — Edit this file, then run: ml launch {filename} -c my-cluster
#
# Full spec: https://docs.skypilot.co/en/latest/reference/yaml-spec.html

resources:
  infra: mithril
  accelerators: B200:8

# Optional: upload this dir (relative to where you run ml launch) to the cluster
# workdir: .

# Optional: runs once when the cluster is first provisioned
# setup: |
#   pip install -r requirements.txt

run: |
  echo "Hello from Mithril!"
"""

# The AGENTS.md index content. {docs_root} is resolved at runtime to the
# absolute path of the bundled docs inside the installed package.
_AGENTS_MD_TEMPLATE = """\
[Mithril CLI Docs]|root: {docs_root}
|Mithril CLI (`ml`) runs ML workloads on GPU clusters.
|CLI: `ml` or `mithril` (interchangeable)
|IMPORTANT: Prefer retrieval-led reasoning over pre-training-led reasoning. \
Read the relevant docs before answering questions about Mithril CLI.
|01-getting-started:{{overview.md,setup.md,quickstart.md}}
|02-launch:{{overview.md}}
|02-launch/concepts:{{autostop.md,detach.md,dryrun.md,gpu-selection.md,\
multi-node.md,spot-instances.md,tasks-clusters-jobs.md}}
|02-launch/yaml:{{file-mounts.md,overview.md,resources.md,task-spec.md}}
|03-clusters:{{overview.md,status.md,exec.md,logs.md,queue.md,stop.md,\
start.md,down.md,ssh.md}}
|04-instances:{{overview.md,create.md,list.md,delete.md,info.md,\
list-types.md,states.md,spot-bidding.md}}
|05-kubernetes:{{overview.md,list.md,info.md,ssh.md,update-kubeconfig.md}}
|06-volumes:{{overview.md,persistent.md,ephemeral.md,file-mounts.md,sync.md}}
|08-sdk:{{overview.md,sky-client.md,mithril-client.md,api-bindings.md,\
launching.md,task-resources.md}}
|09-reference:{{overview.md}}
|09-reference/cli:{{ml.md,ml-launch.md,ml-instance.md,ml-instance-create.md,\
ml-instance-list.md,ml-instance-delete.md,ml-instance-info.md,ml-k8s.md,\
ml-k8s-list.md,ml-k8s-info.md,ml-k8s-ssh.md,ml-k8s-update-kubeconfig.md,\
ml-ssh.md,ml-setup.md}}
|09-reference/api:{{instances-api.md,volumes-api.md,kubernetes-api.md,\
pricing-api.md,models.md}}
|
"""

_AGENT_PROMPT = (
    "Help me run a GPU workload with the Mithril CLI (`ml`)."
    " A task template named {filename} has been generated as a starting point."
)


def agent_docs_root() -> str:
    """Absolute path to the bundled agent docs directory."""
    return str(files("mithril") / "agent" / "docs")


def render_agents_md() -> str:
    """Render the AGENTS.md content with the resolved docs root path."""
    return _AGENTS_MD_TEMPLATE.format(docs_root=agent_docs_root())


def interactive_onboarding(console: Console) -> None:
    """Run the interactive onboarding flow.

    This is the shared entry point called by both `ml launch` (bare invocation)
    and `ml setup` (post-credential configuration).
    """
    # Only run interactive prompts when connected to a terminal.
    # In non-interactive contexts, point to --help and exit.
    if not sys.stdin.isatty():
        console.print(
            "Run [bold]ml launch --help[/bold] for usage.",
            style="dim",
        )
        raise SystemExit(1)

    task_yaml_path = _prompt_and_write_template(console)
    if task_yaml_path is None:
        return

    agents_md_created = _prompt_and_write_agents_md(console)

    if agents_md_created:
        _print_agent_walkthrough(console, task_yaml_path)


def _prompt_and_write_template(console: Console) -> Path | None:
    """Prompt for a task filename and write the template.

    Returns the path that was written, or None if the user cancelled.
    """
    console.print()
    console.print(
        "  [bold]ml launch[/bold] provisions a GPU cluster and runs your workload.\n"
        "  To get started, let's create a task YAML — a file that defines\n"
        "  what to run and what resources you need.\n",
    )

    path = _prompt_template_filename(console)
    if path is None:
        return None

    _write_template(path)
    console.print(
        f"  [green]✓[/green] Created [bold]{path}[/bold]"
        f" — edit this file, then run:"
        f" [bold]ml launch {path} -c my-cluster[/bold]\n",
    )
    return path


def _prompt_template_filename(console: Console) -> Path | None:
    """Pick a filename for the task template.

    Proposes task.yaml, or task-N.yaml if earlier names are taken.
    Returns None on decline or KeyboardInterrupt.
    """
    path = _next_available_path()
    try:
        if not Confirm.ask(
            f"  Create [bold]{path}[/bold]?",
            default=True,
            console=console,
        ):
            return None
    except KeyboardInterrupt:
        console.print()
        return None
    else:
        return path


def _next_available_path() -> Path:
    """Return task.yaml, or task-2.yaml, task-3.yaml, … for the first free name."""
    candidate = Path("task.yaml")
    if not candidate.exists():
        return candidate
    n = 2
    while True:
        candidate = Path(f"task-{n}.yaml")
        if not candidate.exists():
            return candidate
        n += 1


def _write_template(path: Path) -> None:
    """Write the task YAML template to disk."""
    path.write_text(TASK_TEMPLATE.format(filename=path.name))


def _prompt_and_write_agents_md(console: Console) -> bool:
    """Prompt to create/append AGENTS.md. Returns True if it was written."""
    agents_path = Path("AGENTS.md")
    content = render_agents_md()

    try:
        already_exists = agents_path.exists()
        if already_exists:
            existing = agents_path.read_text()
            if "[Mithril CLI Docs]" in existing:
                console.print(
                    "  [dim]AGENTS.md already contains Mithril CLI"
                    " section — skipping.[/dim]\n",
                )
                return True

        if already_exists:
            console.print("  [dim]Detected existing AGENTS.md.[/dim]\n")
            question = (
                "  Add a Mithril CLI section to your AGENTS.md? It points Claude,"
                " Cursor, and other\n"
                "  coding agents to the docs bundled with this"
                " package, so they can\n"
                "  help you write and debug tasks."
            )
        else:
            question = (
                "  Add an AGENTS.md to this project? It points Claude,"
                " Cursor, and other\n"
                "  coding agents to the Mithril CLI docs bundled with"
                " this package,\n"
                "  so they can help you write and debug tasks."
            )

        if not Confirm.ask(question, default=True, console=console):
            return False

        if already_exists:
            with agents_path.open("a") as f:
                f.write("\n" + content)
            console.print(
                "  [green]✓[/green] Appended Mithril CLI section to"
                " [bold]AGENTS.md[/bold]\n",
            )
        else:
            agents_path.write_text(content)
            console.print(
                "  [green]✓[/green] Created [bold]AGENTS.md[/bold]\n",
            )
    except KeyboardInterrupt:
        console.print()
        return False

    return True


def _print_agent_walkthrough(
    console: Console,
    task_yaml_path: Path,
) -> None:
    """Print the agent walkthrough prompt."""
    prompt = _AGENT_PROMPT.format(filename=task_yaml_path.name)
    console.print("  To start an agent-guided walkthrough, use this prompt:\n")
    console.print(f'    [bold]"{prompt}"[/bold]\n')
