"""
Omium CLI tool for local testing and workflow management.
"""

import asyncio
import click
import json
import sys
import httpx
import logging
import getpass
import platform
from typing import Optional, Dict, Any

from omium import OmiumClient
from omium.client import CheckpointError
from omium.config import get_config_manager, ConfigManager, OmiumConfig
from omium.remote_client import RemoteOmiumClient

logger = logging.getLogger(__name__)


def _get_secure_input(prompt: str, default: str = "", allow_visible: bool = False) -> str:
    """
    Get secure input that works reliably on all platforms (Linux, macOS, Windows).
    
    Uses getpass.getpass() which is the standard library solution for cross-platform
    secure input. On Windows, if paste doesn't work, falls back to visible input.
    
    Args:
        prompt: Prompt text to display
        default: Default value if user presses Enter without input
        allow_visible: If True, on Windows will offer visible input as fallback
        
    Returns:
        User input or default value
    """
    is_windows = platform.system() == "Windows"
    
    try:
        # Show prompt first (getpass doesn't show prompts)
        if default:
            click.echo(f"{prompt} [{default[:20]}...] (press Enter to use default, or paste new): ", nl=False)
        else:
            click.echo(f"{prompt}: ", nl=False)
        
        # Use getpass.getpass() - works reliably on all platforms (Linux, macOS, Windows)
        # It handles paste operations better than click.prompt with hide_input
        # Empty string as prompt because we already printed it above
        value = getpass.getpass("")
        
        # Handle empty input - return default if provided
        if not value or not value.strip():
            if default:
                return default
            # On Windows, if empty and allow_visible, offer visible input
            if is_windows and allow_visible and not default:
                click.echo()
                click.echo("💡 Paste didn't work? Trying visible input (your key will be shown):", err=True)
                try:
                    value = click.prompt(prompt, default="", show_default=False)
                except (KeyboardInterrupt, EOFError):
                    raise
                if not value or not value.strip():
                    return ""
        
        return value.strip()
        
    except (KeyboardInterrupt, EOFError):
        # User cancelled input
        raise
    except Exception as e:
        logger.warning(f"Error with getpass, falling back to visible input: {e}")
        # Fallback to visible input if getpass fails
        if allow_visible:
            try:
                click.echo()
                click.echo("⚠️  Secure input failed. Using visible input (your key will be shown):", err=True)
                return click.prompt(prompt, default=default, show_default=bool(default))
            except (KeyboardInterrupt, EOFError):
                raise
        else:
            raise


@click.group()
@click.version_option(version="0.1.3")
def cli():
    """Omium CLI - Fault-tolerant agent operating system."""
    pass


@cli.command("configure")
@click.option("--api-key", help="API key for authentication")
@click.option("--api-url", help="Omium API base URL")
@click.option("--region", help="AWS region")
@click.option("--interactive", "-i", is_flag=True, help="Interactive configuration")
def configure(api_key: Optional[str], api_url: Optional[str], region: Optional[str], interactive: bool):
    """
    Configure Omium SDK settings.
    
    This command sets up your API key, region, and other SDK settings.
    Configuration is saved to ~/.omium/config.json
    """
    config_manager = get_config_manager()
    current_config = config_manager.load()
    
    if interactive:
        click.echo("🔧 Omium SDK Configuration")
        click.echo("=" * 50)
        
        # API Key
        if not api_key:
            system = platform.system()
            if system == "Windows":
                paste_hint = "Paste with Ctrl+V or right-click"
            elif system == "Darwin":  # macOS
                paste_hint = "Paste with Cmd+V"
            else:  # Linux and others
                paste_hint = "Paste with Ctrl+Shift+V or right-click"
            
            click.echo(f"Enter your API key ({paste_hint}):")
            api_key = _get_secure_input(
                "API Key",
                current_config.api_key or "",
            )
            if not api_key:
                click.echo("⚠️  Warning: No API key provided. Set OMIUM_API_KEY environment variable.", err=True)
        
        # API URL
        if not api_url:
            api_url = click.prompt(
                "API URL",
                default=current_config.api_url or "https://api.omium.ai",
            )
        
        # Region
        if not region:
            region = click.prompt(
                "Region",
                default=current_config.region or "us-east-1",
            )
        
        # Validate API key if provided
        if api_key:
            try:
                test_client = RemoteOmiumClient(api_key=api_key, api_url=api_url or current_config.api_url)
                click.echo("✓ API key format is valid")
            except Exception as e:
                click.echo(f"⚠️  Warning: {e}", err=True)
    
    # Update configuration
    updates = {}
    if api_key:
        updates["api_key"] = api_key
    if api_url:
        updates["api_url"] = api_url
    if region:
        updates["region"] = region
    
    if updates:
        config_manager.update(**updates)
        config_manager.save()
        click.echo(f"✓ Configuration saved to {config_manager.config_file}")
    else:
        click.echo("No changes to save. Use --interactive or provide options.")


@cli.command("init")
@click.option("--api-key", help="Omium API key (will prompt if not provided)")
@click.option("--api-url", help="Omium API base URL (defaults to https://api.omium.ai)")
@click.option("--skip-verify", is_flag=True, help="Skip API key verification (not recommended)")
def init(api_key: Optional[str], api_url: Optional[str], skip_verify: bool):
    """
    Initialize Omium SDK configuration.
    
    This command sets up your local Omium environment by:
    1. Asking for your API key (or using --api-key)
    2. Verifying the key with Omium API
    3. Fetching your tenant info and credit balance
    4. Saving configuration to ~/.omium/config.json
    
    After initialization, you can use omium commands without specifying API keys.
    
    Examples:
        # Interactive mode (will prompt for API key)
        omium init
        
        # Non-interactive mode (recommended for Windows/PowerShell)
        omium init --api-key omium_xxxxx --api-url https://api.omium.ai
        
        # Skip verification (not recommended)
        omium init --api-key omium_xxxxx --skip-verify
    
    Note: On Windows/PowerShell, use --api-key flag for best compatibility.
    """
    config_manager = get_config_manager()
    current_config = config_manager.load()
    
    click.echo("🚀 Omium SDK Initialization")
    click.echo("=" * 50)
    click.echo()
    
    # Get API key
    if not api_key:
        click.echo("To get started, you'll need an API key from your Omium account.")
        click.echo("You can create one at: https://app.omium.ai/api-keys")
        click.echo()
        
        # Cross-platform paste instructions
        system = platform.system()
        if system == "Windows":
            paste_hint = "Paste with Ctrl+V or right-click"
        elif system == "Darwin":  # macOS
            paste_hint = "Paste with Cmd+V"
        else:  # Linux and others
            paste_hint = "Paste with Ctrl+Shift+V or right-click"
        
        click.echo(f"💡 Tip: You can paste your API key ({paste_hint})")
        click.echo("   Or use --api-key flag: omium init --api-key YOUR_KEY")
        click.echo()
        
        default_key = current_config.api_key or ""
        if default_key:
            click.echo(f"Found existing API key: {default_key[:20]}...")
            use_existing = click.confirm("Use existing API key?", default=True)
            if use_existing:
                api_key = default_key
            else:
                try:
                    click.echo()
                    click.echo(f"Enter your API key ({paste_hint}):")
                    api_key = _get_secure_input("API Key", "", allow_visible=True)
                except (KeyboardInterrupt, EOFError):
                    click.echo("\n\n❌ Input cancelled.", err=True)
                    click.echo("\nTip: Use --api-key flag to avoid interactive prompts:", err=True)
                    click.echo("   omium init --api-key omium_xxxxx --api-url https://api.omium.ai", err=True)
                    sys.exit(1)
        else:
            try:
                click.echo(f"Enter your API key ({paste_hint}):")
                api_key = _get_secure_input("API Key", "", allow_visible=True)
            except (KeyboardInterrupt, EOFError):
                click.echo("\n\n❌ Input cancelled.", err=True)
                click.echo("\nTip: Use --api-key flag to avoid interactive prompts:", err=True)
                click.echo("   omium init --api-key omium_xxxxx --api-url https://api.omium.ai", err=True)
                sys.exit(1)
    
    if not api_key or not api_key.strip():
        click.echo("❌ API key is required. Exiting.", err=True)
        click.echo("\nGet your API key from: https://app.omium.ai/api-keys", err=True)
        click.echo("\nTip: Use --api-key flag:", err=True)
        click.echo("   omium init --api-key omium_xxxxx --api-url https://api.omium.ai", err=True)
        sys.exit(1)
    
    # Validate API key format
    if not api_key.startswith("omium_"):
        click.echo("⚠️  Warning: API key should start with 'omium_'. Continuing anyway...", err=True)
    
    # Get API URL
    if not api_url:
        default_url = current_config.api_url or "https://api.omium.ai"
        click.echo()
        try:
            api_url = click.prompt(
                "Omium API URL",
                default=default_url,
                show_default=True,
            )
            # Handle empty input
            if not api_url or api_url.strip() == "":
                api_url = default_url
        except (KeyboardInterrupt, EOFError):
            click.echo("\n\nUsing default API URL: https://api.omium.ai")
            api_url = default_url
    
    # Verify API key with Omium API
    if not skip_verify:
        click.echo()
        click.echo("Verifying API key with Omium...")
        
        try:
            tenant_info = asyncio.run(_verify_api_key_and_get_info(api_key, api_url))
            
            if tenant_info:
                click.echo("✓ API key verified successfully!")
                click.echo()
                click.echo(f"Tenant: {tenant_info.get('tenant_name', 'N/A')}")
                click.echo(f"Credits: {tenant_info.get('credits_balance', 0):,} ({tenant_info.get('balance_usd', 0):.2f} USD)")
                click.echo(f"Environment: {tenant_info.get('environment', 'production')}")
                
                if tenant_info.get('credits_balance', 0) == 0:
                    click.echo()
                    click.echo("⚠️  Warning: Your account has no credits.", err=True)
                    click.echo("   Top up at: https://app.omium.ai/billing", err=True)
            else:
                click.echo("⚠️  Warning: Could not verify API key. Saving anyway...", err=True)
                click.echo("   The key may still work for API calls.", err=True)
                click.echo("   If verification continues to fail, check:", err=True)
                click.echo("   - Your API key is correct", err=True)
                click.echo("   - The API URL is correct (https://api.omium.ai)", err=True)
                click.echo("   - You have internet connectivity", err=True)
        except Exception as e:
            error_msg = str(e)
            click.echo(f"⚠️  Warning: Could not verify API key: {error_msg}", err=True)
            click.echo("   Saving configuration anyway. You can verify later.", err=True)
            if "401" in error_msg or "403" in error_msg:
                click.echo("   This might indicate an invalid or inactive API key.", err=True)
            elif "404" in error_msg:
                click.echo("   The verification endpoint might not be available.", err=True)
            elif "timeout" in error_msg.lower() or "connection" in error_msg.lower():
                click.echo("   Network issue detected. Check your internet connection.", err=True)
    
    # Save configuration
    try:
        config_manager.update(api_key=api_key, api_url=api_url)
        config_manager.save()
        
        click.echo()
        click.echo(f"✓ Configuration saved to {config_manager.config_file}")
        click.echo()
        click.echo("You're all set! Next steps:")
        click.echo("  1. Create a workflow: omium init-workflow")
        click.echo("  2. Run a workflow: omium run workflow.json")
        click.echo("  3. View executions: omium list")
        click.echo()
        click.echo("For more help, visit: https://docs.omium.ai")
        
    except Exception as e:
        click.echo(f"❌ Error saving configuration: {e}", err=True)
        sys.exit(1)


async def _verify_api_key_and_get_info(api_key: str, api_url: str) -> Optional[Dict[str, Any]]:
    """
    Verify API key and get tenant information.
    
    Returns:
        Dict with tenant_name, credits_balance, balance_usd, environment, or None if verification fails
    """
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            # Use public SDK verification endpoint
            verify_url = f"{api_url}/api/v1/api-keys/verify"
            logger.debug(f"Verifying API key at: {verify_url}")
            
            verify_response = await client.get(
                verify_url,
                headers={"X-API-Key": api_key},
            )
            
            if verify_response.status_code != 200:
                error_text = verify_response.text[:200] if verify_response.text else "Unknown error"
                error_detail = f"HTTP {verify_response.status_code}: {error_text}"
                logger.warning(f"API key verification failed: {error_detail}")
                # Raise exception with details for better error messages
                raise Exception(error_detail)
            
            key_info = verify_response.json()
            tenant_id = key_info.get("tenant_id")
            
            # Get tenant info and credits
            tenant_info = {
                "tenant_id": tenant_id,
                "tenant_name": key_info.get("tenant_name", "My Account"),
                "environment": key_info.get("environment", "production"),
            }
            
            # Try to get credit balance from billing service
            try:
                balance_response = await client.get(
                    f"{api_url}/api/v1/billing/balance",
                    headers={"X-API-Key": api_key},
                )
                if balance_response.status_code == 200:
                    balance_data = balance_response.json()
                    tenant_info["credits_balance"] = balance_data.get("credits_balance", 0)
                    tenant_info["balance_usd"] = balance_data.get("balance_usd", 0.0)
                else:
                    tenant_info["credits_balance"] = 0
                    tenant_info["balance_usd"] = 0.0
            except Exception as e:
                # Billing service might not be available, that's okay
                logger.debug(f"Could not fetch balance (non-critical): {e}")
                tenant_info["credits_balance"] = 0
                tenant_info["balance_usd"] = 0.0
            
            return tenant_info
            
    except httpx.RequestError as e:
        error_msg = f"Network error: {str(e)}"
        logger.warning(f"Network error verifying API key: {error_msg}")
        raise Exception(error_msg)
    except Exception as e:
        # Re-raise to preserve error details
        if "HTTP" in str(e) or "Network error" in str(e):
            raise
        error_msg = f"Verification error: {str(e)}"
        logger.warning(f"Error verifying API key: {error_msg}")
        raise Exception(error_msg)


@cli.command("init-workflow")
@click.option("--type", type=click.Choice(["crewai", "langgraph"], case_sensitive=False), default="crewai", help="Workflow type")
@click.option("--name", default="my-workflow", help="Workflow name")
@click.option("--output", default="workflow.json", help="Output file path")
def init_workflow(type: str, name: str, output: str):
    """
    Initialize a new workflow template.
    
    Creates a template workflow file that you can customize.
    """
    templates = {
        "crewai": {
            "type": "crewai",
            "workflow_id": name,
            "agent_id": f"{name}-agent",
            "inputs": {
                "topic": "Your input here"
            },
            "definition": {
                "name": name,
                "verbose": True,
                "llm": {
                    "provider": "openai",
                    "model": "gpt-3.5-turbo",
                    "temperature": 0.7
                },
                "agents": [
                    {
                        "role": "Agent",
                        "goal": "Your agent's goal",
                        "backstory": "Your agent's backstory",
                        "verbose": True,
                        "allow_delegation": False
                    }
                ],
                "tasks": [
                    {
                        "description": "Your task description: {topic}",
                        "agent_index": 0,
                        "expected_output": "Expected output description"
                    }
                ]
            }
        },
        "langgraph": {
            "type": "langgraph",
            "workflow_id": name,
            "agent_id": f"{name}-agent",
            "inputs": {
                "messages": [],
                "data": {
                    "input": "Your input here"
                }
            },
            "definition": {
                "name": name,
                "nodes": [
                    {
                        "name": "process",
                        "function": "default_process"
                    },
                    {
                        "name": "respond",
                        "function": "default_respond"
                    }
                ],
                "edges": [
                    {
                        "from": "START",
                        "to": "process"
                    },
                    {
                        "from": "process",
                        "to": "respond"
                    },
                    {
                        "from": "respond",
                        "to": "END"
                    }
                ]
            }
        }
    }
    
    template = templates.get(type.lower(), templates["crewai"])
    
    try:
        with open(output, "w") as f:
            json.dump(template, f, indent=2)
        
        click.echo(f"✓ Created {type} workflow template: {output}")
        click.echo(f"\nNext steps:")
        click.echo(f"  1. Edit {output} to customize your workflow")
        click.echo(f"  2. Set OPENAI_API_KEY environment variable (for CrewAI)")
        click.echo(f"  3. Run: omium run {output}")
        
    except Exception as e:
        click.echo(f"Error creating template: {e}", err=True)
        sys.exit(1)


@cli.command()
@click.argument("workflow_file", type=click.Path(exists=True))
@click.option("--execution-id", help="Execution ID (auto-generated if not provided)")
@click.option("--execution-engine-url", default="http://localhost:8000", help="Execution Engine API URL")
@click.option("--checkpoint-manager", default="localhost:7001", help="Checkpoint Manager URL (for direct checkpoint operations)")
def run(workflow_file: str, execution_id: Optional[str], execution_engine_url: str, checkpoint_manager: str):
    """
    Run a workflow using the Execution Engine.
    
    Workflow file format (JSON):
    {
        "type": "crewai" | "langgraph",
        "workflow_id": "my-workflow",
        "agent_id": "my-agent",
        "inputs": {...},
        "definition": {
            // CrewAI or LangGraph workflow definition
        }
    }
    """
    click.echo(f"Running workflow: {workflow_file}")
    
    # Load workflow file
    try:
        with open(workflow_file, "r") as f:
            workflow_config = json.load(f)
    except json.JSONDecodeError as e:
        click.echo(f"Error: Invalid JSON in workflow file: {e}", err=True)
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error loading workflow file: {e}", err=True)
        sys.exit(1)
    
    # Validate workflow config
    if "type" not in workflow_config and "definition" not in workflow_config:
        click.echo("Error: Workflow file must contain 'type' and 'definition' fields", err=True)
        sys.exit(1)
    
    # Store execution_engine_url in context for async function
    ctx = click.get_current_context()
    ctx.params['execution_engine_url'] = execution_engine_url
    
    # Run workflow
    asyncio.run(_run_workflow(workflow_config, execution_id, checkpoint_manager))


async def _run_workflow(workflow_config: dict, execution_id: Optional[str], checkpoint_manager: str):
    """Run workflow asynchronously."""
    # Get execution engine URL from environment or use default
    execution_engine_url = click.get_current_context().params.get(
        'execution_engine_url',
        'http://localhost:8000'
    )
    
    # Determine workflow type
    workflow_type = workflow_config.get("type", "crewai").lower()
    if workflow_type not in ["crewai", "langgraph"]:
        click.echo(f"Error: Unsupported workflow type '{workflow_type}'. Supported: 'crewai', 'langgraph'", err=True)
        sys.exit(1)
    
    workflow_definition = workflow_config.get("definition", workflow_config)
    inputs = workflow_config.get("inputs", {})
    
    try:
        # Create execution via Execution Engine API
        async with httpx.AsyncClient(timeout=300.0) as http_client:
            # Create execution
            create_response = await http_client.post(
                f"{execution_engine_url}/api/v1/executions",
                json={
                    "workflow_id": workflow_config.get("workflow_id", "cli-workflow"),
                    "agent_id": workflow_config.get("agent_id", "cli-agent"),
                    "input_data": inputs,
                    "metadata": {
                        "workflow_type": workflow_type,
                        "workflow_definition": workflow_definition,
                        "source": "cli"
                    }
                }
            )
            
            if create_response.status_code != 201:
                click.echo(f"Error creating execution: {create_response.text}", err=True)
                sys.exit(1)
            
            execution = create_response.json()
            execution_id = execution["id"]
            
            click.echo(f"✓ Execution created: {execution_id}")
            click.echo(f"✓ Workflow type: {workflow_type}")
            click.echo("Executing workflow...")
            
            # Poll for execution status
            max_wait = 300  # 5 minutes
            wait_time = 0
            poll_interval = 2  # seconds
            
            while wait_time < max_wait:
                await asyncio.sleep(poll_interval)
                wait_time += poll_interval
                
                status_response = await http_client.get(
                    f"{execution_engine_url}/api/v1/executions/{execution_id}"
                )
                
                if status_response.status_code != 200:
                    click.echo(f"Error checking execution status: {status_response.text}", err=True)
                    sys.exit(1)
                
                execution = status_response.json()
                status = execution["status"]
                
                if status == "completed":
                    click.echo("✓ Workflow completed successfully!")
                    click.echo(f"\nResult:\n{json.dumps(execution.get('output_data', {}), indent=2)}")
                    
                    # Show checkpoints if available
                    checkpoints = execution.get("metadata", {}).get("checkpoints", [])
                    if checkpoints:
                        click.echo(f"\n✓ Created {len(checkpoints)} checkpoints")
                        for cp in checkpoints[:5]:  # Show first 5
                            click.echo(f"  - {cp.get('name', 'unknown')}")
                        if len(checkpoints) > 5:
                            click.echo(f"  ... and {len(checkpoints) - 5} more")
                    
                    # Show validation status
                    validation_passed = execution.get("metadata", {}).get("validation_passed")
                    if validation_passed:
                        click.echo("\n✓ Post-condition validation passed")
                    
                    return
                elif status == "failed":
                    error_msg = execution.get("error_message", "Unknown error")
                    click.echo(f"\n✗ Workflow failed: {error_msg}", err=True)
                    
                    # Show failure details
                    metadata = execution.get("metadata", {})
                    failure = metadata.get("failure", {})
                    recovery = metadata.get("recovery", {})
                    suggested_fix = metadata.get("suggested_fix")
                    
                    if failure:
                        click.echo(f"\nFailure Details:")
                        click.echo(f"  Type: {failure.get('type', 'unknown')}")
                        click.echo(f"  Message: {failure.get('message', 'Unknown')}")
                        
                        validation_errors = failure.get("validation_errors", [])
                        if validation_errors:
                            click.echo(f"  Validation Errors:")
                            for err in validation_errors:
                                click.echo(f"    - {err}")
                    
                    if recovery:
                        rollback_performed = recovery.get("rollback_performed", False)
                        if rollback_performed:
                            checkpoint = recovery.get("rollback_checkpoint", {})
                            click.echo(f"\n✓ Automatic rollback performed")
                            click.echo(f"  Rolled back to checkpoint: {checkpoint.get('name', 'unknown')}")
                        else:
                            click.echo(f"\n⚠ No checkpoint available for rollback")
                    
                    if suggested_fix:
                        click.echo(f"\n💡 Suggested Fix:")
                        click.echo(f"  {suggested_fix}")
                        click.echo(f"\nTo retry with fix, update your workflow and run again.")
                    
                    sys.exit(1)
                elif status == "running":
                    click.echo(".", nl=False)  # Progress indicator
                    sys.stdout.flush()
            
            # Timeout
            click.echo(f"\n⚠ Execution timed out after {max_wait} seconds", err=True)
            click.echo(f"Execution ID: {execution_id}")
            click.echo(f"Check status: {execution_engine_url}/api/v1/executions/{execution_id}")
            sys.exit(1)
            
    except httpx.RequestError as e:
        click.echo(f"Error connecting to Execution Engine at {execution_engine_url}: {e}", err=True)
        click.echo("Make sure Execution Engine is running and accessible.", err=True)
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error executing workflow: {e}", err=True)
        sys.exit(1)


@cli.command()
@click.argument("execution_id")
@click.option("--execution-engine-url", default="http://localhost:8000", help="Execution Engine API URL")
@click.option("--checkpoint-manager", default="localhost:7001", help="Checkpoint Manager URL")
@click.option("--checkpoint-id", help="Specific checkpoint ID to replay from (optional)")
def replay(execution_id: str, execution_engine_url: str, checkpoint_manager: str, checkpoint_id: Optional[str]):
    """
    Replay an execution from a checkpoint.
    
    This will recreate the execution state from a previous checkpoint and
    re-execute the workflow from that point.
    """
    click.echo(f"Replaying execution: {execution_id}")
    if checkpoint_id:
        click.echo(f"From checkpoint: {checkpoint_id}")
    asyncio.run(_replay_execution(execution_id, execution_engine_url, checkpoint_manager, checkpoint_id))


async def _replay_execution(
    execution_id: str,
    execution_engine_url: str,
    checkpoint_manager: str,
    checkpoint_id: Optional[str]
):
    """Replay execution from checkpoint."""
    try:
        async with httpx.AsyncClient(timeout=300.0) as http_client:
            # Get original execution
            execution_response = await http_client.get(
                f"{execution_engine_url}/api/v1/executions/{execution_id}"
            )
            
            if execution_response.status_code != 200:
                click.echo(f"Error: Execution {execution_id} not found", err=True)
                sys.exit(1)
            
            execution = execution_response.json()
            
            # Get checkpoints from execution metadata first (they might be stored locally)
            checkpoints = execution.get("metadata", {}).get("checkpoints", [])
            
            # If no checkpoints in metadata, try checkpoint manager
            if not checkpoints:
                click.echo("No checkpoints in execution metadata, trying checkpoint manager...")
                client = OmiumClient(checkpoint_manager_url=checkpoint_manager)
                try:
                    await client.connect()
                    checkpoints = await client.list_checkpoints(execution_id=execution_id)
                    await client.close()
                except Exception as e:
                    click.echo(f"⚠️  Could not get checkpoints from checkpoint manager: {e}", err=True)
                    click.echo("   Checkpoints may be stored locally in execution metadata.", err=True)
                    checkpoints = []
            
            if not checkpoints:
                click.echo("No checkpoints found for this execution", err=True)
                click.echo("Note: Checkpoints may not have been created, or checkpoint manager is unavailable.", err=True)
                sys.exit(1)
            
            click.echo(f"Found {len(checkpoints)} checkpoint(s):")
            for i, cp in enumerate(checkpoints):
                cp_name = cp.get("name") or cp.get("checkpoint_name", "unknown")
                cp_id = cp.get("id", "N/A")
                marker = " <-- Selected" if checkpoint_id and (cp_id == checkpoint_id or cp_id.startswith(checkpoint_id)) else ""
                click.echo(f"  {i+1}. {cp_name} ({cp_id[:16]}...){marker}")
            
            # Select checkpoint
            target_checkpoint = None
            if checkpoint_id:
                target_checkpoint = next(
                    (cp for cp in checkpoints 
                     if cp.get("id") == checkpoint_id or cp.get("id", "").startswith(checkpoint_id)),
                    None
                )
                if not target_checkpoint:
                    click.echo(f"Error: Checkpoint {checkpoint_id} not found", err=True)
                    sys.exit(1)
            else:
                # Use last checkpoint
                target_checkpoint = checkpoints[-1]
                click.echo(f"\nUsing last checkpoint: {target_checkpoint.get('name') or target_checkpoint.get('checkpoint_name', 'unknown')}")
            
            # Use the replay API endpoint (simpler and more reliable)
            click.echo("\nStarting replay via API...")
            replay_response = await http_client.post(
                f"{execution_engine_url}/api/v1/executions/{execution_id}/replay",
                json={"checkpoint_id": target_checkpoint.get("id")} if target_checkpoint.get("id") else {}
            )
            
            if replay_response.status_code == 200:
                result = replay_response.json()
                new_execution_id = result.get("replay_execution_id")
                
                click.echo("✓ Replay started successfully!")
                click.echo(f"\nReplay Execution ID: {new_execution_id}")
                click.echo(f"Original Execution: {execution_id}")
                click.echo(f"Replayed from checkpoint: {target_checkpoint.get('name') or target_checkpoint.get('checkpoint_name', 'unknown')}")
                click.echo(f"\nThe replay execution is running in the background.")
                click.echo(f"Check status: {execution_engine_url}/api/v1/executions/{new_execution_id}")
                click.echo(f"Or use: omium show {new_execution_id}")
            else:
                click.echo(f"Error starting replay: {replay_response.status_code} - {replay_response.text}", err=True)
                sys.exit(1)
                
    except httpx.RequestError as e:
        click.echo(f"Error connecting to services: {e}", err=True)
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error replaying execution: {e}", err=True)
        import traceback
        traceback.print_exc()
        sys.exit(1)


@cli.command()
@click.option("--execution-engine-url", default="http://localhost:8000", help="Execution Engine API URL")
@click.option("--status", help="Filter by status (completed, failed, running, pending)")
@click.option("--workflow-id", help="Filter by workflow ID")
@click.option("--limit", default=20, help="Number of executions to show")
def list(execution_engine_url: str, status: Optional[str], workflow_id: Optional[str], limit: int):
    """
    List recent executions.
    
    Shows a list of recent workflow executions with their status and basic info.
    """
    asyncio.run(_list_executions(execution_engine_url, status, workflow_id, limit))


async def _list_executions(execution_engine_url: str, status: Optional[str], workflow_id: Optional[str], limit: int):
    """List executions asynchronously."""
    try:
        async with httpx.AsyncClient() as client:
            params = {"page": 1, "page_size": limit}
            if status:
                params["status"] = status
            if workflow_id:
                params["workflow_id"] = workflow_id
            
            response = await client.get(
                f"{execution_engine_url}/api/v1/executions",
                params=params
            )
            
            if response.status_code == 200:
                data = response.json()
                executions = data.get("executions", [])
                total = data.get("total", 0)
                
                if not executions:
                    click.echo("No executions found.")
                    return
                
                click.echo(f"\nFound {total} execution(s):\n")
                click.echo(f"{'ID':<40} {'Status':<12} {'Workflow':<30} {'Created'}")
                click.echo("-" * 100)
                
                for exec in executions:
                    exec_id = exec.get("id", "unknown")[:36] + "..."
                    exec_status = exec.get("status", "unknown")
                    workflow = exec.get("workflow_id", "unknown")[:28]
                    created = exec.get("created_at", "unknown")
                    if created != "unknown":
                        from datetime import datetime
                        try:
                            dt = datetime.fromisoformat(created.replace('Z', '+00:00'))
                            created = dt.strftime("%Y-%m-%d %H:%M")
                        except:
                            pass
                    
                    # Color code status
                    if exec_status == "completed":
                        status_display = click.style(exec_status, fg="green")
                    elif exec_status == "failed":
                        status_display = click.style(exec_status, fg="red")
                    elif exec_status == "running":
                        status_display = click.style(exec_status, fg="blue")
                    else:
                        status_display = exec_status
                    
                    click.echo(f"{exec_id:<40} {status_display:<20} {workflow:<30} {created}")
                
                click.echo(f"\nUse 'omium show <execution-id>' to view details")
            else:
                click.echo(f"Error: {response.status_code} - {response.text}", err=True)
                sys.exit(1)
                
    except httpx.RequestError as e:
        click.echo(f"Error connecting to Execution Engine: {e}", err=True)
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error listing executions: {e}", err=True)
        sys.exit(1)


@cli.command()
@click.argument("execution_id")
@click.option("--execution-engine-url", default="http://localhost:8000", help="Execution Engine API URL")
def show(execution_id: str, execution_engine_url: str):
    """
    Show detailed information about an execution.
    
    Displays full execution details including inputs, outputs, checkpoints, and metadata.
    """
    asyncio.run(_show_execution(execution_id, execution_engine_url))


async def _show_execution(execution_id: str, execution_engine_url: str):
    """Show execution details asynchronously."""
    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(
                f"{execution_engine_url}/api/v1/executions/{execution_id}"
            )
            
            if response.status_code == 200:
                execution = response.json()
                
                click.echo(f"\n{'='*60}")
                click.echo(f"Execution: {execution_id}")
                click.echo(f"{'='*60}\n")
                
                # Basic info
                click.echo(f"Status: {execution.get('status', 'unknown')}")
                click.echo(f"Workflow ID: {execution.get('workflow_id', 'unknown')}")
                click.echo(f"Agent ID: {execution.get('agent_id', 'unknown')}")
                click.echo(f"Created: {execution.get('created_at', 'unknown')}")
                click.echo(f"Started: {execution.get('started_at', 'N/A')}")
                click.echo(f"Completed: {execution.get('completed_at', 'N/A')}")
                
                # Input data
                if execution.get("input_data"):
                    click.echo(f"\nInput Data:")
                    click.echo(json.dumps(execution.get("input_data"), indent=2))
                
                # Output data
                if execution.get("output_data"):
                    click.echo(f"\nOutput Data:")
                    click.echo(json.dumps(execution.get("output_data"), indent=2))
                
                # Error message
                if execution.get("error_message"):
                    click.echo(f"\nError:")
                    click.echo(click.style(execution.get("error_message"), fg="red"))
                
                # Checkpoints
                checkpoints = execution.get("metadata", {}).get("checkpoints", [])
                if checkpoints:
                    click.echo(f"\nCheckpoints ({len(checkpoints)}):")
                    for i, cp in enumerate(checkpoints, 1):
                        click.echo(f"  {i}. {cp.get('name', 'unknown')} ({cp.get('id', 'N/A')[:8]}...)")
                        if cp.get("created_at"):
                            click.echo(f"     Created: {cp.get('created_at')}")
                
                # Recovery details
                recovery = execution.get("metadata", {}).get("recovery_details")
                if recovery:
                    click.echo(f"\nRecovery Details:")
                    click.echo(f"  Action: {recovery.get('action', 'N/A')}")
                    click.echo(f"  Status: {recovery.get('status', 'N/A')}")
                    click.echo(f"  Message: {recovery.get('message', 'N/A')}")
                
                click.echo(f"\n{'='*60}")
                click.echo(f"\nUse 'omium replay {execution_id}' to replay this execution")
                
            elif response.status_code == 404:
                click.echo(f"Execution not found: {execution_id}", err=True)
                sys.exit(1)
            else:
                click.echo(f"Error: {response.status_code} - {response.text}", err=True)
                sys.exit(1)
                
    except httpx.RequestError as e:
        click.echo(f"Error connecting to Execution Engine: {e}", err=True)
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error showing execution: {e}", err=True)
        sys.exit(1)


@cli.group()
def checkpoints():
    """Checkpoint management commands."""
    pass


@checkpoints.command("list")
@click.argument("execution_id")
@click.option("--checkpoint-manager", default="localhost:7001", help="Checkpoint Manager URL")
def list_checkpoints(execution_id: str, checkpoint_manager: str):
    """List checkpoints for an execution."""
    asyncio.run(_list_checkpoints(execution_id, checkpoint_manager))


async def _list_checkpoints(execution_id: str, checkpoint_manager: str):
    """List checkpoints asynchronously."""
    client = OmiumClient(checkpoint_manager_url=checkpoint_manager)
    
    try:
        await client.connect()
        
        checkpoints = await client.list_checkpoints(execution_id=execution_id)
        
        if not checkpoints:
            click.echo(f"No checkpoints found for execution {execution_id}")
            click.echo("Note: Checkpoints may be stored locally if checkpoint manager is unavailable.")
            return
        
        click.echo(f"Checkpoints for execution {execution_id}:")
        for cp in checkpoints:
            click.echo(f"  - {cp.get('checkpoint_name', 'unknown')} ({cp.get('id', 'unknown')})")
            if cp.get("created_at"):
                click.echo(f"    Created: {cp['created_at']}")
        
    except Exception as e:
        error_msg = str(e)
        if "connection" in error_msg.lower() or "unavailable" in error_msg.lower():
            click.echo(f"⚠️  Could not connect to checkpoint manager at {checkpoint_manager}", err=True)
            click.echo("   Checkpoints may be stored locally. Check execution metadata for checkpoint info.", err=True)
        else:
            click.echo(f"Error listing checkpoints: {error_msg}", err=True)
        # Don't exit with error code - this is informational
        return
    finally:
        await client.close()


@cli.command()
@click.argument("execution_id")
@click.argument("checkpoint_id")
@click.option("--checkpoint-manager", default="localhost:7001", help="Checkpoint Manager URL")
def rollback(execution_id: str, checkpoint_id: str, checkpoint_manager: str):
    """Rollback execution to a checkpoint."""
    click.echo(f"Rolling back execution {execution_id} to checkpoint {checkpoint_id}")
    asyncio.run(_rollback_execution(execution_id, checkpoint_id, checkpoint_manager))


async def _rollback_execution(execution_id: str, checkpoint_id: str, checkpoint_manager: str):
    """Rollback execution asynchronously."""
    client = OmiumClient(checkpoint_manager_url=checkpoint_manager)
    
    try:
        await client.connect()
        
        result = await client.rollback_to_checkpoint(
            checkpoint_id=checkpoint_id,
            execution_id=execution_id,
            trigger_reason="Manual rollback via CLI",
            trigger_type="manual",
        )
        
        click.echo(f"Rollback successful: {json.dumps(result, indent=2)}")
        
    except CheckpointError as e:
        click.echo(f"Rollback failed: {e}", err=True)
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error during rollback: {e}", err=True)
        sys.exit(1)
    finally:
        await client.close()


@cli.command("export-crew")
@click.argument("source")
@click.option("--output", "-o", help="Output file path (defaults to stdout)")
@click.option("--workflow-name", help="Name for the workflow (defaults to crew name)")
@click.option("--workflow-id", help="ID for the workflow (defaults to workflow-name)")
def export_crew(source: str, output: Optional[str], workflow_name: Optional[str], workflow_id: Optional[str]):
    """
    Export a CrewAI Crew object to Omium workflow format.

    Args:
        source: Path to Python file and object name, e.g., 'path/to/file.py:crew_object'

    Example:
        omium export-crew my_workflow.py:my_crew > workflow.json
        omium export-crew my_workflow.py:my_crew -o workflow.json
    """
    try:
        from omium.adapters.crewai_adapter import export_crewai_workflow
        
        # Parse source path and object name
        if ":" not in source:
            click.echo("Error: Source must be in format 'path/to/file.py:object_name'", err=True)
            sys.exit(1)
        
        file_path, object_name = source.rsplit(":", 1)
        
        # Load the Python file and extract the crew object
        crew = _load_python_object(file_path, object_name)
        
        if crew is None:
            click.echo(f"Error: Could not find object '{object_name}' in {file_path}", err=True)
            sys.exit(1)
        
        # Export the workflow
        workflow = export_crewai_workflow(
            crew=crew,
            workflow_name=workflow_name,
            workflow_id=workflow_id
        )
        
        # Output JSON
        json_output = json.dumps(workflow, indent=2)
        
        if output:
            with open(output, "w") as f:
                f.write(json_output)
            click.echo(f"✓ Exported workflow to {output}")
        else:
            click.echo(json_output)
            
    except ImportError as e:
        click.echo(f"Error: {e}", err=True)
        click.echo("Install CrewAI with: pip install crewai", err=True)
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error exporting CrewAI workflow: {e}", err=True)
        import traceback
        traceback.print_exc()
        sys.exit(1)


@cli.command("export-langgraph")
@click.argument("source")
@click.option("--output", "-o", help="Output file path (defaults to stdout)")
@click.option("--workflow-name", help="Name for the workflow (defaults to graph name)")
@click.option("--workflow-id", help="ID for the workflow (defaults to workflow-name)")
def export_langgraph(source: str, output: Optional[str], workflow_name: Optional[str], workflow_id: Optional[str]):
    """
    Export a LangGraph StateGraph to Omium workflow format.

    Args:
        source: Path to Python file and object name, e.g., 'path/to/file.py:graph_object'

    Note: The graph must be compiled (graph.compile()) before export.

    Example:
        omium export-langgraph my_workflow.py:compiled_graph > workflow.json
        omium export-langgraph my_workflow.py:compiled_graph -o workflow.json
    """
    try:
        from omium.adapters.langgraph_adapter import export_langgraph_workflow
        
        # Parse source path and object name
        if ":" not in source:
            click.echo("Error: Source must be in format 'path/to/file.py:object_name'", err=True)
            sys.exit(1)
        
        file_path, object_name = source.rsplit(":", 1)
        
        # Load the Python file and extract the graph object
        graph = _load_python_object(file_path, object_name)
        
        if graph is None:
            click.echo(f"Error: Could not find object '{object_name}' in {file_path}", err=True)
            sys.exit(1)
        
        # Export the workflow
        workflow = export_langgraph_workflow(
            graph=graph,
            workflow_name=workflow_name,
            workflow_id=workflow_id
        )
        
        # Output JSON
        json_output = json.dumps(workflow, indent=2)
        
        if output:
            with open(output, "w") as f:
                f.write(json_output)
            click.echo(f"✓ Exported workflow to {output}")
        else:
            click.echo(json_output)
            
    except ImportError as e:
        click.echo(f"Error: {e}", err=True)
        click.echo("Install LangGraph with: pip install langgraph", err=True)
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error exporting LangGraph workflow: {e}", err=True)
        import traceback
        traceback.print_exc()
        sys.exit(1)


def _load_python_object(file_path: str, object_name: str):
    """
    Load a Python object from a file.
    
    Args:
        file_path: Path to Python file
        object_name: Name of the object to extract
        
    Returns:
        The loaded object or None if not found
    """
    import importlib.util
    import os
    
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File not found: {file_path}")
    
    # Load the module
    spec = importlib.util.spec_from_file_location("workflow_module", file_path)
    if spec is None or spec.loader is None:
        raise ValueError(f"Could not load module from {file_path}")
    
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    
    # Extract the object
    if not hasattr(module, object_name):
        return None
    
    return getattr(module, object_name)


if __name__ == "__main__":
    cli()
