"""
Cube Agent: Claude SDK agent for EdgescaleAI Cube operations.

Uses the cube-mcp server tools to:
- Connect to Cube clusters via Teleport
- Check Cube status and health
- Deploy Helm charts to Apollo Container Registry

Usage:
    python cube_agent.py "Connect me to staging-int"
    python cube_agent.py --interactive
    python cube_agent.py --setup  # Run setup wizard
"""

import os
import sys
import json
import argparse
import subprocess
import shutil
from pathlib import Path
from typing import Optional

import anthropic
from anthropic import RateLimitError, APIConnectionError
import time

# Agent configuration
AGENT_DIR = Path(__file__).parent
SKILLS_DIR = AGENT_DIR / ".claude" / "skills"
LLM_CONFIG_PATH = AGENT_DIR / "llm_config.json"

# Default model
MODEL = "claude-sonnet-4-20250514"
MAX_TOKENS = 8192


# ============================================================================
# Tool Prerequisites Check
# ============================================================================

REQUIRED_TOOLS = {
    "tsh": {
        "name": "Teleport Client",
        "check": ["tsh", "version"],
        "install": "https://goteleport.com/download/",
        "description": "Required for Cube cluster authentication"
    },
    "kubectl": {
        "name": "Kubernetes CLI",
        "check": ["kubectl", "version", "--client"],
        "install": "brew install kubectl",
        "description": "Required for Cube node operations"
    },
    "helm": {
        "name": "Helm",
        "check": ["helm", "version", "--short"],
        "install": "brew install helm",
        "description": "Required for Helm chart deployments"
    },
    "docker": {
        "name": "Docker",
        "check": ["docker", "--version"],
        "install": "https://docs.docker.com/get-docker/",
        "description": "Required for building container images"
    },
    "apollo-cli": {
        "name": "Apollo CLI",
        "check": ["which", "apollo-cli"],
        "install": "Download from Apollo portal",
        "description": "Required for publishing to Apollo"
    }
}


def check_tool(tool_cmd: str, check_args: list) -> tuple[bool, str]:
    """Check if a tool is installed and return version info."""
    try:
        result = subprocess.run(
            check_args,
            capture_output=True,
            text=True,
            timeout=10
        )
        if result.returncode == 0:
            version = result.stdout.strip() or result.stderr.strip()
            # For 'which' command, just show "Installed" instead of path
            if check_args[0] == "which":
                return True, "Installed"
            return True, version.split('\n')[0][:50]
        return False, "Command failed"
    except FileNotFoundError:
        return False, "Not found"
    except subprocess.TimeoutExpired:
        return False, "Timeout"
    except Exception as e:
        return False, str(e)


def check_all_tools() -> dict:
    """Check all required tools and return status."""
    status = {}
    for tool, info in REQUIRED_TOOLS.items():
        installed, version = check_tool(tool, info["check"])
        status[tool] = {
            "installed": installed,
            "version": version,
            "name": info["name"],
            "install": info["install"],
            "description": info["description"]
        }
    return status


def check_docker_buildx() -> tuple[bool, str]:
    """Check if Docker buildx is available."""
    try:
        result = subprocess.run(
            ["docker", "buildx", "version"],
            capture_output=True,
            text=True,
            timeout=10
        )
        if result.returncode == 0:
            return True, result.stdout.strip().split('\n')[0][:50]
        return False, "Not available"
    except:
        return False, "Not available"


def check_apollo_credentials() -> tuple[bool, str, str]:
    """Check if Apollo credentials are set."""
    client_id = os.environ.get("APOLLO_CLIENT")
    client_secret = os.environ.get("APOLLO_SECRET")

    if client_id and client_secret:
        return True, client_id[:8] + "...", "Set"
    elif client_id:
        return False, client_id[:8] + "...", "APOLLO_SECRET missing"
    elif client_secret:
        return False, "Missing", "APOLLO_CLIENT missing"
    else:
        return False, "Missing", "Both APOLLO_CLIENT and APOLLO_SECRET missing"


def print_status_report():
    """Print a comprehensive status report of all tools and config."""
    print("\n" + "=" * 60)
    print("CUBE AGENT - ENVIRONMENT STATUS")
    print("=" * 60)

    # Check tools
    print("\nRequired Tools:")
    print("-" * 60)
    tool_status = check_all_tools()
    all_tools_ok = True

    for tool, info in tool_status.items():
        status = "OK" if info["installed"] else "MISSING"
        icon = "[+]" if info["installed"] else "[ ]"
        print(f"  {icon} {info['name']:<20} {status:<10} {info['version']}")
        if not info["installed"]:
            all_tools_ok = False
            print(f"      Install: {info['install']}")

    # Check Docker buildx
    buildx_ok, buildx_ver = check_docker_buildx()
    icon = "[+]" if buildx_ok else "[ ]"
    print(f"  {icon} {'Docker Buildx':<20} {'OK' if buildx_ok else 'MISSING':<10} {buildx_ver}")
    if not buildx_ok and tool_status.get("docker", {}).get("installed"):
        print("      Install: docker buildx install")

    # Check Apollo credentials
    print("\nApollo Container Registry:")
    print("-" * 60)
    creds_ok, client_info, secret_info = check_apollo_credentials()
    icon = "[+]" if creds_ok else "[ ]"
    print(f"  {icon} APOLLO_CLIENT:  {client_info}")
    print(f"  {icon} APOLLO_SECRET:  {secret_info}")

    if not creds_ok:
        print("\n  To configure Apollo credentials, add to ~/.bashrc or ~/.zshrc:")
        print('    export APOLLO_CLIENT="<your-client-id>"')
        print('    export APOLLO_SECRET="<your-client-secret>"')
        print("  Then run: source ~/.bashrc")

    # Check LLM config
    print("\nLLM Configuration:")
    print("-" * 60)
    if LLM_CONFIG_PATH.exists():
        try:
            config = json.loads(LLM_CONFIG_PATH.read_text())
            provider = config.get("provider", "anthropic")
            print(f"  [+] Provider: {provider}")
            if provider == "anthropic_bedrock":
                print(f"  [+] Model: {config.get('anthropic_bedrock', {}).get('model', 'default')}")
            else:
                print(f"  [+] Model: {config.get('anthropic', {}).get('model', 'default')}")
        except:
            print("  [ ] Error reading llm_config.json")
    else:
        api_key = os.environ.get("ANTHROPIC_API_KEY")
        if api_key:
            print(f"  [+] ANTHROPIC_API_KEY: Set ({api_key[:8]}...)")
        else:
            print("  [ ] ANTHROPIC_API_KEY: Not set")
            print("      Set with: export ANTHROPIC_API_KEY='your-key'")
            print("      Or copy llm_config.example.json to llm_config.json")

    print("\n" + "=" * 60)

    # Summary
    if all_tools_ok and creds_ok:
        print("STATUS: Ready to use!")
    elif all_tools_ok:
        print("STATUS: Tools OK, but Apollo credentials missing (Helm deploy won't work)")
    else:
        print("STATUS: Some tools missing. Install them to enable full functionality.")

    print("=" * 60 + "\n")

    return all_tools_ok


# ============================================================================
# LLM Client Setup
# ============================================================================

def load_llm_config() -> dict:
    """Load LLM configuration from llm_config.json if it exists."""
    if LLM_CONFIG_PATH.exists():
        try:
            return json.loads(LLM_CONFIG_PATH.read_text())
        except Exception as e:
            print(f"Warning: Failed to load llm_config.json: {e}")
    return {}


def create_client():
    """Create the appropriate Anthropic client based on config."""
    config = load_llm_config()
    provider = config.get("provider", "anthropic_bedrock")  # Default to Bedrock

    if provider == "anthropic_bedrock":
        bedrock_cfg = config.get("anthropic_bedrock", {})
        region = bedrock_cfg.get("aws_region") or os.environ.get("AWS_REGION", "us-west-2")

        # Check if explicit credentials are provided in config or env
        access_key = bedrock_cfg.get("aws_access_key_id") or os.environ.get("AWS_ACCESS_KEY_ID")
        secret_key = bedrock_cfg.get("aws_secret_access_key") or os.environ.get("AWS_SECRET_ACCESS_KEY")

        if access_key and secret_key:
            # Use explicit credentials
            return anthropic.AnthropicBedrock(
                aws_access_key=access_key,
                aws_secret_key=secret_key,
                aws_region=region,
            ), bedrock_cfg.get("model", MODEL)
        else:
            # Use default AWS credential chain (~/.aws/credentials, SSO, IAM role, etc.)
            return anthropic.AnthropicBedrock(
                aws_region=region,
            ), bedrock_cfg.get("model", MODEL)
    else:
        # Direct Anthropic API
        anthropic_cfg = config.get("anthropic", {})
        api_key = anthropic_cfg.get("api_key") or os.environ.get("ANTHROPIC_API_KEY")

        if not api_key:
            raise ValueError("ANTHROPIC_API_KEY not set. Export it or configure llm_config.json")

        return anthropic.Anthropic(api_key=api_key), anthropic_cfg.get("model", MODEL)


# ============================================================================
# System Prompt
# ============================================================================

def load_skills() -> str:
    """Load all skills files into a single context string."""
    skills_content = []

    if SKILLS_DIR.exists():
        for skill_dir in sorted(SKILLS_DIR.iterdir()):
            if skill_dir.is_dir():
                skill_file = skill_dir / "SKILL.md"
                if skill_file.exists():
                    content = skill_file.read_text()
                    if content.startswith("---"):
                        parts = content.split("---", 2)
                        if len(parts) >= 3:
                            content = parts[2].strip()
                    skills_content.append(f"## Skill: {skill_dir.name}\n\n{content}")

    return "\n\n---\n\n".join(skills_content)


def build_system_prompt() -> str:
    """Build the system prompt for the Cube agent."""

    skills = load_skills()

    # Get current tool status for context
    tool_status = check_all_tools()
    creds_ok, _, _ = check_apollo_credentials()

    tool_summary = []
    for tool, info in tool_status.items():
        status = "installed" if info["installed"] else "MISSING"
        tool_summary.append(f"- {info['name']}: {status}")

    system_prompt = f"""You are the Cube Agent, an assistant for managing EdgescaleAI Cube edge devices and deployments.

## Your Capabilities

You can help users with:

1. **Cube Cluster Access** (via Teleport)
   - Login to Cube clusters (staging-int, staging-pre-prod, ironmountain-v4, etc.)
   - List available Cubes
   - Check Cube node status and health

2. **Helm Deployments** (via Apollo Container Registry)
   - Build and push Docker images
   - Package and push Helm charts
   - Publish to Apollo catalog
   - Auto-detect Chart.yaml and version bumping

## Current Environment Status

Tools:
{chr(10).join(tool_summary)}

Apollo Credentials: {'Configured' if creds_ok else 'NOT CONFIGURED'}

## Available Tools

### cube_login
Connect to a Cube cluster via Teleport SSO.
- Opens browser for authentication
- Sets up kubectl access automatically

### cube_list
List all available Cube clusters you can connect to.

### cube_status
Get detailed status of the connected Cube node (conditions, capacity, resources).

### acr_login
Authenticate to Apollo Container Registry (Docker + Helm registries).
Requires APOLLO_CLIENT and APOLLO_SECRET environment variables.

### acr_get_token
Get an OAuth2 token from Apollo (for debugging).

### helm_deploy
Full deployment pipeline:
1. Login to ACR
2. Build & push Docker image (if Dockerfile exists)
3. Bump Chart.yaml version
4. Package Helm chart
5. Push to ACR
6. Publish to Apollo

Options:
- chart_path: Path to chart (auto-detects if not specified)
- version: Exact version (e.g., "1.0.0")
- bump: "patch", "minor", or "major" (default: patch)
- app_version: Docker image tag version
- skip_docker: Skip Docker build
- dry_run: Preview without changes

## Workflow Examples

**Connecting to a Cube:**
```
User: "Connect me to staging-int"
You: Call cube_login(cluster="staging-int")
```

**Checking Cube Health:**
```
User: "What's the status of my Cube?"
You: Call cube_status()
```

**Deploying a Helm Chart:**
```
User: "Deploy my chart"
You: Call helm_deploy() with auto-detection

User: "Deploy as version 2.0.0"
You: Call helm_deploy(version="2.0.0")

User: "Deploy with minor version bump, dry run first"
You: Call helm_deploy(bump="minor", dry_run=True)
```

## Setup Guidance

If a user hasn't configured their environment, help them:

1. **Missing Tools**: Provide installation commands
2. **Missing Apollo Credentials**: Show them how to set APOLLO_CLIENT and APOLLO_SECRET
3. **Missing API Key**: Help them configure ANTHROPIC_API_KEY or llm_config.json

## Important Notes

- Always check prerequisites before attempting operations
- For helm_deploy, offer dry_run first for new users
- Teleport login opens a browser - warn users about this
- Cube operations require being connected first (cube_login)
- ACR operations require Apollo credentials

{skills}
"""

    return system_prompt


# ============================================================================
# Tool Definitions & Handlers
# ============================================================================

def create_tool_definitions() -> list:
    """Define tools that map to cube-mcp server functions."""

    return [
        {
            "name": "cube_login",
            "description": """Login to a Cube cluster via Teleport (opens browser for SSO).
This will authenticate you and set up kubectl access to the Cube.""",
            "input_schema": {
                "type": "object",
                "properties": {
                    "cluster": {
                        "type": "string",
                        "description": "Cube cluster name (e.g., staging-int, staging-pre-prod, ironmountain-v4)"
                    },
                    "proxy": {
                        "type": "string",
                        "description": "Teleport proxy URL (default: edgescaleai.teleport.sh)"
                    }
                },
                "required": ["cluster"]
            }
        },
        {
            "name": "cube_list",
            "description": "List all available Cube clusters accessible via Teleport.",
            "input_schema": {
                "type": "object",
                "properties": {},
                "required": []
            }
        },
        {
            "name": "cube_status",
            "description": "Get detailed status of the connected Cube node (conditions, capacity, resources).",
            "input_schema": {
                "type": "object",
                "properties": {},
                "required": []
            }
        },
        {
            "name": "acr_get_token",
            "description": "Get an OAuth2 access token from Apollo Container Registry. Useful for debugging.",
            "input_schema": {
                "type": "object",
                "properties": {},
                "required": []
            }
        },
        {
            "name": "acr_login",
            "description": "Login to Apollo Container Registry for both Docker and Helm. Requires APOLLO_CLIENT and APOLLO_SECRET.",
            "input_schema": {
                "type": "object",
                "properties": {},
                "required": []
            }
        },
        {
            "name": "helm_deploy",
            "description": """Full Helm deployment pipeline: build Docker image, package chart, push to ACR, publish to Apollo.

Auto-detects Chart.yaml and Dockerfile. Use dry_run=true to preview first.""",
            "input_schema": {
                "type": "object",
                "properties": {
                    "chart_path": {
                        "type": "string",
                        "description": "Path to chart directory or Chart.yaml (auto-detected if not specified)"
                    },
                    "version": {
                        "type": "string",
                        "description": "Exact version to use (e.g., '1.0.0'). If not specified, bumps current version."
                    },
                    "bump": {
                        "type": "string",
                        "enum": ["patch", "minor", "major"],
                        "description": "Version bump type (default: patch). Ignored if version is specified."
                    },
                    "app_version": {
                        "type": "string",
                        "description": "Docker image tag version. If not specified, uses chart version."
                    },
                    "skip_docker": {
                        "type": "boolean",
                        "description": "Skip Docker build even if Dockerfile exists."
                    },
                    "dry_run": {
                        "type": "boolean",
                        "description": "Preview what would be done without making changes."
                    }
                },
                "required": []
            }
        }
    ]


# Import the actual tool implementations from cube-mcp
# We import them here to reuse the logic
sys.path.insert(0, str(AGENT_DIR.parent / "src"))
try:
    from cube_mcp.server import (
        cube_login as _cube_login,
        cube_list as _cube_list,
        cube_status as _cube_status,
        acr_get_token as _acr_get_token,
        acr_login as _acr_login,
        helm_deploy as _helm_deploy
    )
    MCP_AVAILABLE = True
except ImportError:
    MCP_AVAILABLE = False
    print("Warning: cube_mcp module not found. Install with: pip install -e /path/to/cube-mcp")


def handle_tool_call(tool_name: str, tool_input: dict) -> str:
    """Route tool calls to cube-mcp server functions."""

    if not MCP_AVAILABLE:
        return "Error: cube_mcp module not installed. Run: pip install -e /path/to/cube-mcp"

    try:
        if tool_name == "cube_login":
            return _cube_login(
                cluster=tool_input["cluster"],
                proxy=tool_input.get("proxy", "edgescaleai.teleport.sh")
            )

        elif tool_name == "cube_list":
            return _cube_list()

        elif tool_name == "cube_status":
            return _cube_status()

        elif tool_name == "acr_get_token":
            return _acr_get_token()

        elif tool_name == "acr_login":
            return _acr_login()

        elif tool_name == "helm_deploy":
            return _helm_deploy(
                chart_path=tool_input.get("chart_path"),
                version=tool_input.get("version"),
                bump=tool_input.get("bump", "patch"),
                app_version=tool_input.get("app_version"),
                skip_docker=tool_input.get("skip_docker", False),
                dry_run=tool_input.get("dry_run", False)
            )

        else:
            return f"Unknown tool: {tool_name}"

    except Exception as e:
        return f"Error executing {tool_name}: {e}"


# ============================================================================
# Agent Loop
# ============================================================================

def run_agent(prompt: str, max_turns: int = 30) -> str:
    """Run the Cube agent with the given prompt."""

    # Initialize client
    client, model = create_client()
    print(f"Using model: {model}")

    # Build system prompt and tools
    system_prompt = build_system_prompt()
    tools = create_tool_definitions()

    # Initialize conversation
    messages = [
        {"role": "user", "content": prompt}
    ]

    print(f"\n{'=' * 60}")
    print("CUBE AGENT")
    print(f"{'=' * 60}")
    print(f"Prompt: {prompt[:200]}{'...' if len(prompt) > 200 else ''}")
    print(f"{'=' * 60}\n")

    turn = 0
    while turn < max_turns:
        turn += 1
        print(f"\n--- Turn {turn} ---")

        # Call Claude with retry on rate limit
        for attempt in range(5):
            try:
                response = client.messages.create(
                    model=model,
                    max_tokens=MAX_TOKENS,
                    system=system_prompt,
                    tools=tools,
                    messages=messages
                )
                break
            except (RateLimitError, APIConnectionError) as e:
                wait = 2 ** attempt
                print(f"  Rate limited, waiting {wait}s...")
                time.sleep(wait)

        # Check stop reason
        if response.stop_reason == "end_turn":
            final_text = ""
            for block in response.content:
                if hasattr(block, "text"):
                    final_text += block.text
            print(f"\n{'=' * 60}")
            print("RESPONSE")
            print(f"{'=' * 60}")
            print(final_text)
            return final_text

        elif response.stop_reason == "tool_use":
            assistant_content = response.content
            messages.append({"role": "assistant", "content": assistant_content})

            tool_results = []
            for block in assistant_content:
                if block.type == "tool_use":
                    print(f"  Tool: {block.name}")
                    print(f"    Input: {json.dumps(block.input)[:100]}...")

                    result = handle_tool_call(block.name, block.input)

                    # Truncate result for display
                    display_result = result[:300] + "..." if len(result) > 300 else result
                    print(f"    Result: {display_result}")

                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": result
                    })

            messages.append({"role": "user", "content": tool_results})

        else:
            print(f"Unexpected stop reason: {response.stop_reason}")
            break

    print(f"\nAgent reached max turns ({max_turns})")
    return "Agent reached maximum turns without completing"


# ============================================================================
# CLI Entry Point
# ============================================================================

def interactive_input(prompt_str: str = "> ") -> str:
    """
    Read input with Enter to submit, Ctrl+J/Ctrl+Enter for newline.
    """
    try:
        from prompt_toolkit import prompt
        from prompt_toolkit.key_binding import KeyBindings
        from prompt_toolkit.keys import Keys

        bindings = KeyBindings()

        @bindings.add(Keys.ControlJ)
        def _(event):
            """Ctrl+J inserts a newline."""
            event.current_buffer.insert_text('\n')

        return prompt(prompt_str, multiline=True, key_bindings=bindings)

    except ImportError:
        # Fallback to simple input
        return input(prompt_str)


def interactive_mode(max_turns: int = 30, session_id: str = None):
    """Run in interactive chat mode with session persistence."""
    # Import session management
    sys.path.insert(0, str(AGENT_DIR))
    from app.sessions import get_or_create_session, list_sessions, SESSIONS_DIR

    print("\n" + "=" * 60)
    print("CUBE AGENT - Interactive Mode")
    print("=" * 60)
    print("Enter: Submit | Ctrl+J: New line | Ctrl+C: Exit")
    print("Commands: /sessions, /new, /load <id>, /history")
    print(f"Sessions saved to: {SESSIONS_DIR}")
    print("=" * 60)

    # Show status once
    print_status_report()

    # Get or create session
    session = get_or_create_session(session_id)
    print(f"\nSession: {session.session_id}")
    if session.messages:
        print(f"Continuing conversation ({len(session.messages)} messages)")

    while True:
        try:
            prompt = interactive_input("\n> ").strip()

            if not prompt:
                continue

            if prompt.lower() in ["exit", "quit", "q"]:
                print(f"\nSession saved: {session.file_path}")
                print("Goodbye!")
                break

            # Handle commands
            if prompt.startswith("/"):
                cmd_parts = prompt[1:].split(maxsplit=1)
                cmd = cmd_parts[0].lower()

                if cmd == "sessions":
                    sessions = list_sessions()
                    print("\nSaved Sessions:")
                    print("-" * 60)
                    for s in sessions[:10]:
                        msg_count = len(s.messages)
                        print(f"  {s.session_id}  ({msg_count} msgs)  {s.title[:40]}")
                    if not sessions:
                        print("  No sessions found")
                    continue

                elif cmd == "new":
                    session = get_or_create_session()
                    print(f"\nNew session: {session.session_id}")
                    continue

                elif cmd == "load" and len(cmd_parts) > 1:
                    new_session = get_or_create_session(cmd_parts[1])
                    if new_session:
                        session = new_session
                        print(f"\nLoaded session: {session.session_id}")
                        print(f"Messages: {len(session.messages)}")
                    else:
                        print(f"Session not found: {cmd_parts[1]}")
                    continue

                elif cmd == "history":
                    print(f"\nSession: {session.session_id}")
                    print(f"File: {session.file_path}")
                    print("-" * 60)
                    for msg in session.messages[-10:]:
                        role = msg["role"]
                        content = msg["content"]
                        if isinstance(content, str):
                            preview = content[:100] + "..." if len(content) > 100 else content
                            print(f"  [{role}] {preview}")
                        else:
                            print(f"  [{role}] <{len(content)} blocks>")
                    continue

                elif cmd == "help":
                    print("\nCommands:")
                    print("  /sessions     List saved sessions")
                    print("  /new          Start new session")
                    print("  /load <id>    Load a session")
                    print("  /history      Show recent messages")
                    print("  /help         Show this help")
                    continue

            # Add user message to session
            session.add_user_message(prompt)

            # Run agent
            result = run_agent_with_session(prompt, session, max_turns=max_turns)

            print()  # Extra newline after response

        except KeyboardInterrupt:
            print(f"\n\nSession saved: {session.file_path}")
            print("Goodbye!")
            break
        except Exception as e:
            print(f"\nError: {e}")
            import traceback
            traceback.print_exc()


def run_agent_with_session(prompt: str, session, max_turns: int = 30) -> str:
    """Run the agent using an existing session for history."""

    # Initialize client
    client, model = create_client()

    # Build system prompt and tools
    system_prompt = build_system_prompt()
    tools = create_tool_definitions()

    print(f"\n{'=' * 60}")
    print(f"Processing: {prompt[:50]}...")
    print(f"{'=' * 60}\n")

    turn = 0
    while turn < max_turns:
        turn += 1
        print(f"--- Turn {turn} ---")

        # Call Claude with retry
        for attempt in range(5):
            try:
                response = client.messages.create(
                    model=model,
                    max_tokens=MAX_TOKENS,
                    system=system_prompt,
                    tools=tools,
                    messages=session.get_messages_for_api()
                )
                break
            except (RateLimitError, APIConnectionError) as e:
                wait = 2 ** attempt
                print(f"  Rate limited, waiting {wait}s...")
                time.sleep(wait)

        # Check stop reason
        if response.stop_reason == "end_turn":
            final_text = ""
            for block in response.content:
                if hasattr(block, "text"):
                    final_text += block.text

            session.add_assistant_message(response.content)

            print(f"\n{'=' * 60}")
            print("RESPONSE")
            print(f"{'=' * 60}")
            print(final_text)
            return final_text

        elif response.stop_reason == "tool_use":
            assistant_content = response.content
            session.add_assistant_message(assistant_content)

            tool_results = []
            for block in assistant_content:
                if block.type == "tool_use":
                    print(f"  Tool: {block.name}")
                    print(f"    Input: {json.dumps(block.input)[:100]}...")

                    result = handle_tool_call(block.name, block.input)

                    display_result = result[:300] + "..." if len(result) > 300 else result
                    print(f"    Result: {display_result}")

                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": result
                    })

            session.add_tool_results(tool_results)

        else:
            print(f"Unexpected stop reason: {response.stop_reason}")
            break

    return "Agent reached maximum turns"


def main():
    parser = argparse.ArgumentParser(
        description="Cube Agent: Manage EdgescaleAI Cube devices and deployments"
    )
    parser.add_argument(
        "prompt",
        nargs="?",
        help="The prompt/command for the agent"
    )
    parser.add_argument(
        "--interactive", "-i",
        action="store_true",
        help="Run in interactive mode"
    )
    parser.add_argument(
        "--setup",
        action="store_true",
        help="Run interactive setup wizard"
    )
    parser.add_argument(
        "--session",
        help="Session ID to continue (or 'new' for fresh session)"
    )
    parser.add_argument(
        "--sessions",
        action="store_true",
        help="List saved sessions"
    )
    parser.add_argument(
        "--status", "-s",
        action="store_true",
        help="Show environment status and exit"
    )
    parser.add_argument(
        "--skip-auth-check",
        action="store_true",
        help="Skip authentication check on startup"
    )
    parser.add_argument(
        "--max-turns",
        type=int,
        default=30,
        help="Maximum conversation turns (default 30)"
    )

    args = parser.parse_args()

    # Status check mode
    if args.status:
        print_status_report()
        sys.exit(0)

    # Setup wizard
    if args.setup:
        sys.path.insert(0, str(AGENT_DIR))
        from setup.wizard import run_setup_wizard
        run_setup_wizard(force=True)
        sys.exit(0)

    # Check auth before running (unless skipped)
    if not args.skip_auth_check and not args.sessions:
        sys.path.insert(0, str(AGENT_DIR))
        from setup.wizard import check_and_setup
        config, auth_status = check_and_setup()

    # List sessions
    if args.sessions:
        sys.path.insert(0, str(AGENT_DIR))
        from app.sessions import list_sessions, SESSIONS_DIR
        print(f"\nSaved Sessions ({SESSIONS_DIR}):")
        print("-" * 60)
        for s in list_sessions():
            msg_count = len(s.messages)
            print(f"  {s.session_id}  ({msg_count} msgs)  {s.title[:40]}")
        sys.exit(0)

    # Interactive mode
    if args.interactive:
        session_id = None if args.session == "new" else args.session
        interactive_mode(max_turns=args.max_turns, session_id=session_id)
        sys.exit(0)

    # Show status first
    print_status_report()

    # Get prompt
    if args.prompt:
        prompt = args.prompt
    else:
        # Default to a helpful intro
        prompt = "What can you help me with? Give me a brief overview of your capabilities."

    # Run agent (with session if specified)
    if args.session:
        sys.path.insert(0, str(AGENT_DIR))
        from app.sessions import get_or_create_session
        session_id = None if args.session == "new" else args.session
        session = get_or_create_session(session_id)
        session.add_user_message(prompt)
        run_agent_with_session(prompt.strip(), session, max_turns=args.max_turns)
        print(f"\nSession saved: {session.file_path}")
    else:
        run_agent(prompt.strip(), max_turns=args.max_turns)


if __name__ == "__main__":
    main()
