🚀 Snowflake Setup CLI with Copilot Agent - Implementation Plan

#!/usr/bin/env python3
"""
Snowflake Setup CLI - Automated Snowflake configuration with Copilot Agent.

This CLI guides users through:
1. Generating PAT (Programmatic Access Token) SQL
2. Creating Snow CLI connection
3. Setting up complete Snowflake infrastructure from SNOWFLAKE_COMPLETE_SETUP.md
"""

import asyncio
from pathlib import Path

import typer
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt
from rich.syntax import Syntax

app = typer.Typer(
    name="snowflake-setup",
    help="Automated Snowflake setup with Copilot Agent assistance",
    add_completion=False,
)
console = Console()


@app.command()
def setup(
    account: str = typer.Option(
        None,
        "--account",
        "-a",
        help="Snowflake account identifier (e.g., <account_identifier>)",
        prompt="Enter your Snowflake account identifier",
    ),
    user: str = typer.Option(
        None,
        "--user",
        "-u",
        help="Snowflake username",
        prompt="Enter your Snowflake username",
    ),
    token_name: str = typer.Option(
        "EVSNOW_SETUP_PAT",
        "--token-name",
        "-t",
        help="Name for the PAT token",
    ),
    days_to_expiry: int = typer.Option(
        90,
        "--days",
        "-d",
        help="Days until PAT expires",
    ),
) -> None:
    """
    Start the interactive Snowflake setup wizard.

    This command will:
    1. Generate SQL to create a PAT for your user
    2. Ask you to run the SQL and provide the token
    3. Start an AI agent to complete the setup automatically
    """
    asyncio.run(_run_setup(account, user, token_name, days_to_expiry))


async def _run_setup(
    account: str,
    user: str,
    token_name: str,
    days_to_expiry: int,
) -> None:
    """Main setup orchestration."""
    from .sql_templates import generate_pat_sql
    from .agent import run_setup_agent

    console.print(Panel.fit(
        "[bold blue]Snowflake Setup Wizard[/bold blue]\n"
        "This wizard will help you set up your complete Snowflake infrastructure.",
        title="🏔️ EvSnow Setup"
    ))

    # Step 1: Generate and display PAT SQL
    pat_sql = generate_pat_sql(user, token_name, days_to_expiry)

    console.print("\n[bold yellow]Step 1: Create Programmatic Access Token (PAT)[/bold yellow]")
    console.print("\nRun the following SQL in your Snowflake worksheet as ACCOUNTADMIN:\n")

    syntax = Syntax(pat_sql, "sql", theme="monokai", line_numbers=True)
    console.print(Panel(syntax, title="SQL to Run", border_style="green"))

    console.print("\n[bold]After running the SQL, copy the 'token_secret' value from the output.[/bold]")
    console.print("[dim]The token appears ONLY ONCE - save it immediately![/dim]\n")

    # Step 2: Get the PAT from user
    pat_token = Prompt.ask(
        "[bold cyan]Paste the token_secret here[/bold cyan]",
        password=True,  # Hide input for security
    )

    if not pat_token or len(pat_token) < 10:
        console.print("[bold red]Error: Invalid token provided. Please run the setup again.[/bold red]")
        raise typer.Exit(1)

    console.print("\n[bold green]✓ Token received![/bold green]")

    # Step 3: Start the Copilot agent
    console.print("\n[bold yellow]Step 2: Starting AI Agent for automated setup...[/bold yellow]\n")

    await run_setup_agent(
        account=account,
        user=user,
        pat_token=pat_token,
    )


@app.command()
def validate() -> None:
    """Validate the current Snowflake setup."""
    console.print("[bold]Validating Snowflake setup...[/bold]")
    # TODO: Implement validation


if __name__ == "__main__":
    app()

"""SQL template generation for Snowflake PAT and setup."""


def generate_pat_sql(
    user: str,
    token_name: str,
    days_to_expiry: int = 90,
) -> str:
    """
    Generate SQL to create a PAT for the specified user.

    Args:
        user: Snowflake username
        token_name: Name for the PAT
        days_to_expiry: Token validity period

    Returns:
        SQL statement to execute as ACCOUNTADMIN
    """
    return f"""\
-- Create Programmatic Access Token (PAT) for automated setup
-- Run this as ACCOUNTADMIN in a Snowflake SQL worksheet

-- First, ensure you're using ACCOUNTADMIN role
USE ROLE ACCOUNTADMIN;

-- Create the PAT
ALTER USER {user} ADD PROGRAMMATIC ACCESS TOKEN {token_name}
    ROLE_RESTRICTION = 'ACCOUNTADMIN'
    DAYS_TO_EXPIRY = {days_to_expiry}
    COMMENT = 'PAT for EvSnow automated setup';

-- IMPORTANT: The token_secret appears ONLY in this output!
-- Copy it immediately and store securely.
-- You cannot retrieve it again later.
"""


def generate_verify_pat_sql(user: str) -> str:
    """Generate SQL to verify PAT tokens for a user."""
    return f"""\
-- List all PATs for user
SHOW USER PROGRAMMATIC ACCESS TOKENS FOR USER {user};
"""

"""System prompts and instructions for the Copilot agent."""

from pathlib import Path


def get_setup_guide_content() -> str:
    """Read the SNOWFLAKE_COMPLETE_SETUP.md file content."""
    guide_path = Path(__file__).parent.parent.parent / "SNOWFLAKE_COMPLETE_SETUP.md"
    if guide_path.exists():
        return guide_path.read_text()
    return ""


def get_agent_system_prompt(account: str, user: str) -> str:
    """
    Generate the system prompt for the setup agent.

    The agent will:
    1. Use Snow CLI to create and test a connection
    2. Follow SNOWFLAKE_COMPLETE_SETUP.md to set up everything
    """
    setup_guide = get_setup_guide_content()

    return f"""\
You are an expert Snowflake setup assistant. Your task is to help the user set up their complete Snowflake infrastructure for the EvSnow project.

## Context
- Snowflake Account: {account}
- Snowflake User: {user}
- The user has already created a PAT (Programmatic Access Token) with ACCOUNTADMIN role restriction

## Your Tasks (Execute in Order)

### Task 1: Create Snow CLI Connection
Use the Snow CLI to create a connection using the PAT token.

1. Create a connection using:
```bash
snow connection add \\
    --connection-name evsnow-setup \\
    --account {account} \\
    --user {user} \\
    --authenticator PROGRAMMATIC_ACCESS_TOKEN \\
    --token-file-path <path-to-token-file>

"""GitHub Copilot SDK agent for Snowflake setup automation."""

import asyncio
import tempfile
from pathlib import Path

from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel

from copilot import CopilotClient
from copilot.generated.session_events import SessionEventType

from .prompts import get_agent_system_prompt

console = Console()


async def run_setup_agent(
    account: str,
    user: str,
    pat_token: str,
) -> None:
    """
    Run the Copilot agent to complete Snowflake setup.

    Args:
        account: Snowflake account identifier
        user: Snowflake username
        pat_token: PAT token for authentication
    """
    # Save PAT token to a temporary file for Snow CLI
    token_file = Path(tempfile.mktemp(suffix=".token"))
    token_file.write_text(pat_token)
    token_file.chmod(0o600)  # Secure permissions

    try:
        # Create Copilot client
        client = CopilotClient()
        await client.start()

        console.print("[dim]Agent started. Processing setup...[/dim]\n")

        # Get system prompt with setup guide
        system_prompt = get_agent_system_prompt(account, user)

        # Create session with the agent
        session = await client.create_session({
            "model": "gpt-4.1",  # Or claude-sonnet-4.5
            "streaming": True,
            "system_message": {
                "mode": "append",
                "content": system_prompt,
            },
            "on_user_input_request": handle_user_input,
        })

        # Set up event handlers for streaming output
        done = asyncio.Event()

        def on_event(event):
            if event.type.value == "assistant.message_delta":
                # Stream the response as it comes
                delta = event.data.delta_content or ""
                console.print(delta, end="")
            elif event.type.value == "assistant.message":
                # Final message
                console.print()  # New line after streaming
            elif event.type.value == "tool.execution_start":
                tool_name = event.data.tool_name if hasattr(event.data, 'tool_name') else 'tool'
                console.print(f"\n[dim]→ Executing: {tool_name}[/dim]")
            elif event.type.value == "tool.execution_complete":
                console.print("[dim]  ✓ Complete[/dim]")
            elif event.type.value == "session.idle":
                done.set()

        session.on(on_event)

        # Send initial prompt to start the setup
        initial_prompt = f"""\
Start the Snowflake setup process.

The PAT token has been saved to: {token_file}

Please begin by:
1. Creating a Snow CLI connection using this token file
2. Testing the connection
3. Then proceeding with the full setup

Account: {account}
User: {user}
Token file: {token_file}
"""

        await session.send({"prompt": initial_prompt})

        # Wait for completion
        await done.wait()

        # Cleanup
        await session.destroy()
        await client.stop()

        console.print(Panel.fit(
            "[bold green]✓ Setup process completed![/bold green]\n\n"
            "Review the output above for any manual steps needed.",
            title="🎉 Complete"
        ))

    finally:
        # Clean up token file
        if token_file.exists():
            token_file.unlink()


async def handle_user_input(request, invocation):
    """
    Handle requests from the agent that need user input.

    This is called when the agent uses the ask_user tool.
    """
    question = request.get("question", "")
    choices = request.get("choices", [])
    allow_freeform = request.get("allowFreeform", True)

    console.print(f"\n[bold cyan]Agent asks:[/bold cyan] {question}")

    if choices:
        console.print("[dim]Options:[/dim]")
        for i, choice in enumerate(choices, 1):
            console.print(f"  {i}. {choice}")

    from rich.prompt import Prompt
    answer = Prompt.ask("[bold cyan]Your answer[/bold cyan]")

    return {
        "answer": answer,
        "wasFreeform": allow_freeform,
    }