"""MLflow CLI commands for Assistant integration."""

import sys
import threading
import time
from pathlib import Path

import click

from mlflow.assistant.config import AssistantConfig, ProjectConfig, SkillsConfig
from mlflow.assistant.providers import AssistantProvider, list_providers
from mlflow.assistant.providers.base import ProviderNotConfiguredError
from mlflow.assistant.skill_installer import install_skills


class Spinner:
    """Simple spinner animation for long-running operations."""

    def __init__(self, message: str = "Loading"):
        self.message = message
        self.spinning = False
        self.thread = None
        self.frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]

    def _spin(self):
        i = 0
        while self.spinning:
            frame = self.frames[i % len(self.frames)]
            sys.stdout.write(f"\r{frame} {self.message}")
            sys.stdout.flush()
            time.sleep(0.1)
            i += 1

    def __enter__(self):
        self.spinning = True
        self.thread = threading.Thread(target=self._spin, name="Spinner")
        self.thread.start()
        return self

    def __exit__(self, *args):
        self.spinning = False
        if self.thread:
            self.thread.join()
        sys.stdout.write("\r" + " " * (len(self.message) + 4) + "\r")
        sys.stdout.flush()


@click.command("assistant")
@click.option(
    "--configure",
    is_flag=True,
    help="Configure or reconfigure the assistant settings",
)
def commands(configure: bool):
    """MLflow Assistant - AI-powered trace analysis.

    Run 'mlflow assistant --configure' to set up the assistant.
    """
    if configure:
        _run_configuration()
    else:
        # Check if already configured
        config = AssistantConfig.load()
        if not config.providers:
            click.secho(
                "Assistant is not configured. Please run: mlflow assistant --configure",
                fg="yellow",
            )
        else:
            click.secho(
                "Assistant launch is not yet implemented. To use Assistant, run `mlflow assistant "
                "--configure` to setup, then launch the MLflow UI manually.",
                fg="yellow",
            )


def _run_configuration():
    """Configure MLflow Assistant for the UI.

    This interactive command sets up the AI assistant feature that allows you
    to analyze MLflow traces directly from the UI.

    The command will:
    1. Ask which provider to use (Claude Code for now)
    2. Check provider availability
    3. Optionally connect an experiment with code repository
    4. Ask which model to use
    5. Ask where to install skills (user-level or project-level)
    6. Install provider-specific skills
    7. Save configuration

    Example:
        mlflow assistant --configure
    """
    click.echo()
    click.secho("╔══════════════════════════════════════════╗", fg="cyan")
    click.secho("║       *    .  *       .   *              ║", fg="cyan")
    click.secho("║   .    *  MLflow Assistant Setup   *  .  ║", fg="cyan", bold=True)
    click.secho("║      *    .       *   .      *           ║", fg="cyan")
    click.secho("╚══════════════════════════════════════════╝", fg="cyan")
    click.echo()

    # Step 1: Select provider
    provider = _prompt_provider()
    if provider is None:
        return

    # Step 2: Check provider availability
    if not _check_provider(provider):
        return

    # Step 3: Optionally connect experiment with code repository
    project_path = _prompt_experiment_path()

    # Step 4: Ask for model
    model = _prompt_model()

    # Step 5: Ask for skill location
    skills_config = _prompt_skill_location(project_path)

    # Step 6: Install skills
    skill_path = _install_skills(provider, skills_config, project_path)

    # Step 7: Save configuration
    _save_config(provider, model, skills_config)

    # Show success message
    _show_init_success(provider, model, skill_path)


def _prompt_provider() -> AssistantProvider | None:
    """Prompt user to select a provider."""
    providers = list_providers()

    click.secho("Step 1/4: Select AI Provider", fg="cyan", bold=True)
    click.secho("-" * 30, fg="cyan")
    click.echo()

    for i, provider in enumerate(providers, 1):
        marker = click.style(" [recommended]", fg="green") if i == 1 else ""
        click.echo(f"  {i}. {provider.display_name}{marker}")
        click.secho(f"     {provider.description}", dim=True)

    click.echo()
    click.secho("  More providers coming soon...", dim=True)
    click.echo()

    default_provider = providers[0]
    choice = click.prompt(
        click.style(f"Select provider [1: {default_provider.display_name}]", fg="bright_blue"),
        default="1",
        type=click.Choice([str(i) for i in range(1, len(providers) + 1)]),
        show_choices=False,
        show_default=False,
    )

    provider = providers[int(choice) - 1]
    click.echo()
    return provider


def _check_provider(provider: AssistantProvider) -> bool:
    click.secho("Step 2/4: Checking Provider", fg="cyan", bold=True)
    click.secho("-" * 30, fg="cyan")
    click.echo()

    if not provider.is_available():
        click.secho(
            f"{provider.display_name} is not available. "
            "Please ensure it is installed and accessible in your PATH.",
            fg="red",
        )
        click.echo()
        return False

    try:
        spinner_msg = "Checking connection... " + click.style(
            "(this may take a few seconds)", dim=True
        )
        with Spinner(spinner_msg):
            provider.check_connection()
        click.secho("Connection verified", fg="green")
        click.echo()
        return True
    except ProviderNotConfiguredError as e:
        click.secho(str(e), fg="red")
        click.echo()
        return False


def _fetch_recent_experiments(tracking_uri: str, max_results: int = 5) -> list[tuple[str, str]]:
    """Fetch recent experiments from the tracking server.

    Returns:
        List of (experiment_id, experiment_name) tuples.
    """
    import mlflow

    original_uri = mlflow.get_tracking_uri()
    try:
        mlflow.set_tracking_uri(tracking_uri)
        client = mlflow.MlflowClient()
        experiments = client.search_experiments(
            max_results=max_results,
            order_by=["last_update_time DESC"],
        )
        return [(exp.experiment_id, exp.name) for exp in experiments]
    except Exception:
        return []
    finally:
        mlflow.set_tracking_uri(original_uri)


def _resolve_experiment_id(tracking_uri: str, name_or_id: str) -> str | None:
    """Resolve experiment name or ID to experiment ID.

    Args:
        tracking_uri: MLflow tracking server URI.
        name_or_id: Experiment name or ID.

    Returns:
        Experiment ID if found, None otherwise.
    """
    import mlflow

    original_uri = mlflow.get_tracking_uri()
    try:
        mlflow.set_tracking_uri(tracking_uri)
        client = mlflow.MlflowClient()

        # First try to get by ID (if it looks like an ID)
        if name_or_id.isdigit():
            try:
                if exp := client.get_experiment(name_or_id):
                    return exp.experiment_id
            except Exception:
                pass

        # Try to get by name
        if exp := client.get_experiment_by_name(name_or_id):
            return exp.experiment_id

        return None
    except Exception:
        return None
    finally:
        mlflow.set_tracking_uri(original_uri)


def _prompt_experiment_path() -> Path | None:
    """Prompt user to optionally connect an experiment with code repository.

    Returns:
        The project path if configured, None otherwise.
    """
    click.secho("Step 3/5: Experiment & Code Context ", fg="cyan", bold=True, nl=False)
    click.secho("[Optional, Recommended]", fg="green", bold=True)
    click.secho("-" * 30, fg="cyan")
    click.echo()
    click.echo("You can connect an experiment with a code repository to give")
    click.echo("the assistant context about your source code for better analysis.")
    click.secho("(You can also set this up later in the MLflow UI.)", dim=True)
    click.echo()

    connect = click.confirm(
        click.style(
            "Do you want to connect an experiment with a code repository?", fg="bright_blue"
        ),
        default=True,
    )

    if not connect:
        click.echo()
        return None

    click.echo()

    # Ask for tracking URI to fetch experiments
    tracking_uri = click.prompt(
        click.style("Enter the MLflow tracking server URI", fg="bright_blue"),
        default="http://localhost:5000",
    )

    click.echo()
    click.secho("Fetching recent experiments...", dim=True)

    # Fetch recent experiments
    experiments = _fetch_recent_experiments(tracking_uri)

    if not experiments:
        click.secho("Could not fetch experiments from the server.", fg="yellow")
        click.echo("You can set this up later in the MLflow UI.")
        click.echo()
        return None

    click.echo()
    click.echo(click.style("Select an experiment to connect:", fg="bright_blue"))
    click.echo()

    for i, (exp_id, exp_name) in enumerate(experiments, 1):
        click.echo(f"  {i}. {exp_name} (ID: {exp_id})")

    other_option = len(experiments) + 1
    click.echo(f"  {other_option}. Enter experiment name or ID manually")
    click.echo()

    choice = click.prompt(
        click.style("Select experiment", fg="bright_blue"),
        type=click.IntRange(1, other_option),
        default=1,
    )

    if choice == other_option:
        while True:
            click.echo()
            name_or_id = click.prompt(
                click.style("Experiment name or ID", fg="bright_blue"), default=""
            )
            if not name_or_id:
                click.secho("No experiment specified. Please try again.", fg="yellow")
                continue

            experiment_id = _resolve_experiment_id(tracking_uri, name_or_id)
            if experiment_id:
                # Use the input as display name (could be name or ID)
                experiment_name = name_or_id
                break

            click.secho(
                f"Experiment '{name_or_id}' not found. Please try again.",
                fg="red",
            )
    else:
        experiment_id, experiment_name = experiments[choice - 1]

    click.secho(
        f"Experiment '{experiment_name}' selected",
        fg="green",
    )
    click.echo()

    # Ask for project path
    default_path = str(Path.cwd())
    while True:
        raw_path = click.prompt(
            click.style("Enter the path to your project directory:", fg="bright_blue"),
            default=default_path,
        )
        # Expand ~ and resolve relative paths
        expanded_path = Path(raw_path).expanduser().resolve()
        if expanded_path.is_dir():
            project_path = str(expanded_path)
            break
        click.secho(f"Directory '{raw_path}' does not exist. Please try again.", fg="red")

    # Save the project path mapping locally
    try:
        config = AssistantConfig.load()
        config.projects[experiment_id] = ProjectConfig(type="local", location=project_path)
        config.save()
        click.secho(
            f"Project path {project_path} is saved for experiment '{experiment_name}'",
            fg="green",
        )
    except Exception as e:
        click.secho(f"Error saving project path: {e}", fg="red")

    click.echo()
    return expanded_path


def _prompt_model() -> str:
    """Prompt user for model selection."""
    click.secho("Step 4/5: Model Selection", fg="cyan", bold=True)
    click.secho("-" * 30, fg="cyan")
    click.echo()
    click.echo("Choose a model for analysis:")
    click.secho("  - Press Enter to use the default model (recommended)", dim=True)
    click.secho("  - Or type a specific model name (e.g., claude-sonnet-4-20250514)", dim=True)
    click.echo()

    model = click.prompt(click.style("Model", fg="bright_blue"), default="default")
    click.echo()
    return model


def _prompt_skill_location(project_path: Path | None) -> SkillsConfig:
    """Prompt user for skill installation location.

    Args:
        project_path: The project path from experiment setup, or None if skipped.

    Returns:
        SkillsConfig with the selected location type and optional custom path.
    """
    click.secho("Step 5/5: Skill Installation Location", fg="cyan", bold=True)
    click.secho("-" * 30, fg="cyan")
    click.echo()
    click.echo("Choose where to install MLflow skills for Assistant:")
    click.echo()

    # TODO: Update this when we support other providers
    user_path = Path.home() / ".claude" / "skills"
    click.echo(f"  1. User level ({user_path})")
    click.secho("     Skills available globally across all projects", dim=True)
    click.echo()

    if project_path:
        project_skill_path = project_path / ".claude" / "skills"
        click.echo(f"  2. Project level ({project_skill_path})")
        click.secho("     Skills available only in this project", dim=True)
        click.echo()
        click.echo("  3. Custom location")
        click.secho("     Specify a custom path for skills", dim=True)
        click.echo()
        valid_choices = ["1", "2", "3"]
    else:
        click.echo("  2. Custom location")
        click.secho("     Specify a custom path for skills", dim=True)
        click.echo()
        valid_choices = ["1", "2"]

    choice = click.prompt(
        click.style("Select location [1: User level]", fg="bright_blue"),
        default="1",
        type=click.Choice(valid_choices),
        show_choices=False,
        show_default=False,
    )

    click.echo()

    if choice == "1":
        return SkillsConfig(type="global")
    elif choice == "2" and project_path:
        return SkillsConfig(type="project")
    else:
        # Custom location
        while True:
            raw_path = click.prompt(
                click.style("Enter the custom path for skills", fg="bright_blue"),
                default=str(user_path),
            )
            expanded_path = Path(raw_path).expanduser().resolve()
            # For custom paths, we'll create the directory, so just check parent exists
            if expanded_path.parent.exists() or expanded_path.exists():
                click.echo()
                return SkillsConfig(type="custom", custom_path=str(expanded_path))
            click.secho(
                f"Parent directory '{expanded_path.parent}' does not exist. Please try again.",
                fg="red",
            )


def _install_skills(
    provider: AssistantProvider, skills_config: SkillsConfig, project_path: Path | None
) -> Path:
    """Install skills bundled with MLflow.

    Returns:
        The resolved path where skills were installed.
    """
    match skills_config.type:
        case "global":
            skill_path = provider.resolve_skills_path(Path.home())
        case "project":
            if project_path is None:
                raise ValueError("project_path is required for 'project' skills location")
            skill_path = provider.resolve_skills_path(project_path)
        case "custom":
            if skills_config.custom_path is None:
                raise ValueError("custom_path is required for 'custom' skills location")
            skill_path = Path(skills_config.custom_path).expanduser()
    if installed_skills := install_skills(skill_path):
        for skill in installed_skills:
            click.secho(f"  - {skill}")
    else:
        click.secho("No skills available to install.", fg="yellow")
    click.echo()
    return skill_path


def _save_config(provider: AssistantProvider, model: str, skills_config: SkillsConfig) -> None:
    """Save configuration to file."""
    click.secho("Saving Configuration", fg="cyan", bold=True)
    click.secho("-" * 30, fg="cyan")

    config = AssistantConfig.load()
    config.set_provider(provider.name, model)
    config.providers[provider.name].skills = skills_config
    config.save()

    click.secho("Configuration saved", fg="green")
    click.echo()


def _show_init_success(provider: AssistantProvider, model: str, skill_path: Path) -> None:
    """Show success message and next steps."""
    click.secho("  ~ * ~ * ~ * ~ * ~ * ~ * ~ * ~", fg="green")
    click.secho("        Setup Complete!        ", fg="green", bold=True)
    click.secho("  ~ * ~ * ~ * ~ * ~ * ~ * ~ * ~", fg="green")
    click.echo()
    click.secho("Configuration:", bold=True)
    click.echo(f"  Provider: {provider.display_name}")
    click.echo(f"  Model: {model}")
    click.echo(f"  Skills: {skill_path}")
    click.echo()
    click.secho("Next steps:", bold=True)
    click.echo("  1. Start MLflow server:")
    click.secho("     $ mlflow server", fg="cyan")
    click.echo()
    click.echo("  2. Open MLflow UI and navigate to an experiment")
    click.echo()
    click.echo("  3. Click 'Ask Assistant'")
    click.echo()
    click.secho("To reconfigure, run: ", nl=False)
    click.secho("mlflow assistant --configure", fg="cyan")