import click
import asyncio
import os
import shutil
from pathlib import Path
from rich.console import Console
from rich.panel import Panel
from .config import RunnerConfig
from .auth import save_token, validate_token
from .utils import setup_logger, get_log_shipper
from .jupyter_backend import JupyterBackend
from .file_manager import FileManager
from .message_handler import MessageHandler
from .client import RunnerClient
from .output_buffer import OutputBuffer
from .notification_multiplexer import NotificationMultiplexer
from .local_server import LocalWebSocketServer
from .compat import setup_signal_handlers
from .state_snapshot import start_snapshot_loop
from .structured_log import get_logger as _get_slog

_venv_log = _get_slog("cloud_venv")

console = Console()


def ensure_cloud_venv():
    """Create or verify /workspace/.venv for persistent package storage.

    Uses --system-site-packages so pre-baked packages (numpy, pandas, etc.)
    are inherited from the Docker image at zero disk cost.
    """
    import subprocess
    import sys

    venv_path = Path("/workspace/.venv")
    venv_python = venv_path / "bin" / "python"

    # Check if venv exists and is functional
    if venv_python.exists():
        try:
            result = subprocess.run(
                [str(venv_python), "--version"],
                capture_output=True, text=True, timeout=5,
            )
            if result.returncode == 0:
                # Check Python version matches system
                venv_ver = ".".join(result.stdout.strip().split()[-1].split(".")[:2])
                sys_result = subprocess.run(
                    [sys.executable, "--version"],
                    capture_output=True, text=True, timeout=5,
                )
                sys_ver = ".".join(sys_result.stdout.strip().split()[-1].split(".")[:2])
                if venv_ver == sys_ver:
                    # Verify ipykernel is installed (may have failed on previous boot)
                    ipy_check = subprocess.run(
                        [str(venv_python), "-c", "import ipykernel"],
                        capture_output=True, timeout=5,
                    )
                    if ipy_check.returncode == 0:
                        _venv_log.info("cloud_venv.exists", python_version=venv_ver)
                        return
                    # ipykernel missing — recreate
                    _venv_log.info("cloud_venv.recreated", reason="ipykernel_missing")
                    shutil.rmtree(venv_path, ignore_errors=True)
                else:
                    # Version mismatch — recreate
                    _venv_log.info("cloud_venv.recreated", old_version=venv_ver, new_version=sys_ver)
                    shutil.rmtree(venv_path, ignore_errors=True)
        except Exception:
            # Broken venv — recreate
            _venv_log.info("cloud_venv.recreated", reason="broken")
            shutil.rmtree(venv_path, ignore_errors=True)

    # Create venv with system-site-packages
    try:
        subprocess.run(
            ["uv", "venv", str(venv_path),
             "--python", sys.executable,
             "--system-site-packages"],
            check=True, timeout=30, capture_output=True,
        )
        subprocess.run(
            ["uv", "pip", "install", "ipykernel",
             "--python", str(venv_python)],
            check=True, timeout=60, capture_output=True,
        )
        _venv_log.info("cloud_venv.created")
    except FileNotFoundError:
        _venv_log.warning("cloud_venv.failed", error="uv not found")
    except subprocess.CalledProcessError as e:
        _venv_log.warning("cloud_venv.failed", error=str(e))
    except Exception as e:
        _venv_log.warning("cloud_venv.failed", error=str(e))


@click.group()
def cli():
    """Skop Runner - Local execution client for Skop Labs"""
    pass


@cli.command()
@click.option("--token", required=True, help="Device token from Skop Labs")
def login(token: str):
    """Authenticate runner with device token."""
    config = RunnerConfig.load()

    console.print("[cyan]Validating token...[/cyan]")

    try:
        # Validate token
        result = asyncio.run(validate_token_with_token(token, config))

        # Save token
        save_token(token, config)

        console.print(
            Panel.fit(
                f"[green]Authentication successful![/green]\n\n"
                f"User: {result['user']['email']}\n"
                f"Runner ID: {result['runnerId']}",
                title="Logged In",
            )
        )

    except Exception as e:
        console.print(f"[red]Authentication failed: {e}[/red]")
        raise click.Abort()


async def validate_token_with_token(token: str, config: RunnerConfig):
    """Helper to validate token."""
    temp_config = RunnerConfig(
        device_token=token,
        api_url=config.api_url,
        ws_url=config.ws_url,
        log_level=config.log_level,
        config_dir=config.config_dir,
    )
    return await validate_token(temp_config)


@cli.command()
def start():
    """Start the runner and connect to Skop API."""
    is_cloud = os.getenv("SKOP_CLOUD_RUNNER") == "true"
    config = RunnerConfig.load()

    if not config.device_token:
        console.print("[red]Not authenticated. Run 'skop-runner login' first.[/red]")
        raise click.Abort()

    # Setup logging (local runners get a LogShipper processor for HTTP shipping)
    setup_logger(config.log_level, api_url=config.api_url)

    # Cloud runner: copy demo template contents into /workspace on first start.
    # .skop/ directory is created later by workspace.setRoot — we rely on that
    # as the marker to skip re-copy on subsequent starts. The inner `not dst.exists()`
    # guard prevents overwriting user files even if the marker is absent.
    if is_cloud:
        demo_src = Path("/templates/demo")
        marker = Path("/workspace/.skop")
        if demo_src.exists() and not marker.exists():
            try:
                for item in demo_src.iterdir():
                    dst = Path("/workspace") / item.name
                    if not dst.exists():
                        if item.is_dir():
                            shutil.copytree(item, dst)
                        else:
                            shutil.copy2(item, dst)
                console.print("[green]Demo content initialized in /workspace[/green]")
            except Exception as e:
                console.print(f"[yellow]Could not copy demo template: {e}[/yellow]")

        # Create persistent venv for package storage
        ensure_cloud_venv()

    from . import __version__
    console.print(
        Panel.fit(
            f"[cyan]Starting Skop Runner v{__version__}[/cyan]\n\n"
            f"API URL: {config.api_url}\n"
            f"WS URL: {config.ws_url}"
            + (f"\nMode: cloud" if is_cloud else ""),
            title="Runner Starting",
        )
    )

    # Initialize components - no project root yet, will be set via workspace.setRoot
    # JupyterBackend encapsulates Jupyter Server lifecycle + kernel management.
    # The server is NOT started eagerly - it waits for workspace.setRoot to know the project root
    file_manager = FileManager(is_cloud=is_cloud)

    # Define callback that notifies frontend when backend dies.
    # message_handler is captured by closure and will be set before the callback
    # can ever fire (backend starts lazily on workspace.setRoot).
    message_handler = None

    def on_backend_died():
        """Notify frontend when kernel backend dies and restarts are exhausted."""
        console.print("[red]Kernel backend died and couldn't be restarted[/red]")

        if message_handler is not None:
            async def send_notification():
                await message_handler.send_notification("kernel.backendDied", {
                    "message": "Kernel backend crashed and automatic restart failed. "
                               "Kernel execution is unavailable until the runner is restarted.",
                })

            # Schedule the async notification from sync callback
            try:
                asyncio.create_task(send_notification())
            except RuntimeError:
                # No event loop running (shouldn't happen in normal operation)
                pass

    def on_kernel_died(kernel_id: str):
        """Clean up session, execution queue, and notify frontend when a kernel dies."""
        console.print(f"[yellow]Kernel {kernel_id} died unexpectedly[/yellow]")

        # Clean up stale session so get_or_create_session doesn't return it
        if message_handler is not None:
            message_handler.session_manager.remove_sessions_for_kernel(kernel_id)

            # Clean up execution queue for dead kernel
            asyncio.create_task(
                message_handler._execution_queue.cleanup_kernel(kernel_id)
            )

            async def send_notification():
                await message_handler.send_notification("kernel.died", {
                    "kernel_id": kernel_id,
                })

            try:
                asyncio.create_task(send_notification())
            except RuntimeError:
                pass

    def on_kernel_status_change(kernel_id: str, execution_state: str):
        """Push kernel status changes to frontend."""
        if message_handler is not None:
            async def send_notification():
                await message_handler.send_notification("kernel.statusChanged", {
                    "kernel_id": kernel_id,
                    "execution_state": execution_state,
                })

            try:
                asyncio.create_task(send_notification())
            except RuntimeError:
                pass

    # Persistence directory for surviving runner restarts
    # Cloud runners use /workspace/.skop/state (persisted via Fly volume)
    if is_cloud:
        persistence_dir = "/workspace/.skop/state"
        Path(persistence_dir).mkdir(parents=True, exist_ok=True)
    else:
        persistence_dir = str(Path.home() / ".skop" / "state")

    jupyter_backend = JupyterBackend(
        on_backend_died=on_backend_died,
        on_kernel_died=on_kernel_died,
        on_backend_restarted=lambda: console.print("[cyan]Kernel backend restarted[/cyan]"),
        on_kernel_status_change=on_kernel_status_change,
        persistence_dir=persistence_dir,
    )
    output_buffer = OutputBuffer()
    message_handler = MessageHandler(
        jupyter_backend, file_manager,
        output_buffer=output_buffer,
        persistence_dir=persistence_dir,
        idle_timeout=2592000 if is_cloud and os.getenv("SKOP_POOL_MACHINE") != "true" else None,  # 30 days idle for user cloud runners, disabled for pool
    )
    client = RunnerClient(config, message_handler, output_buffer)

    first_connect = True

    def on_connect():
        nonlocal first_connect
        if first_connect:
            console.print("[green]Connected to Skop Labs[/green]")
            console.print("[dim]Runner is now active. Press Ctrl+C to stop.[/dim]\n")
            first_connect = False
        else:
            console.print("[green]Reconnected to Skop Labs[/green]")

    client.on_connect = on_connect
    shutdown_event = asyncio.Event()

    # Cloud runners don't need local WS server (no localhost browser access)
    local_server = None if is_cloud else LocalWebSocketServer(message_handler)

    async def run():
        loop = asyncio.get_running_loop()

        # Handle signals for graceful shutdown using asyncio
        def signal_handler():
            console.print("\n[yellow]Shutting down...[/yellow]")
            client.running = False
            shutdown_event.set()
            # Schedule disconnect to cancel blocking tasks
            asyncio.create_task(client.disconnect())

        setup_signal_handlers(loop, signal_handler)

        try:
            # Start log shipper for local runners (ships logs to /api/logs)
            log_shipper = get_log_shipper()
            if log_shipper:
                await log_shipper.start()

            # Initialize backend (creates server manager but doesn't start server).
            # Server starts lazily when workspace.setRoot fires from the frontend.
            # Kernel reconnection + session restore happen at that point too.
            await jupyter_backend.initialize()

            # Start local WS server (optional — relay is the primary path)
            # Skipped for cloud runners (no localhost browser access)
            local_ws_port = None
            if local_server is not None:
                try:
                    local_ws_port = await local_server.start()
                    console.print(f"[dim]Local WS server on port {local_ws_port}[/dim]")
                except Exception as e:
                    console.print(f"[yellow]Local WS server failed to start: {e}[/yellow]")
                    console.print("[dim]Continuing with relay-only mode[/dim]")

            # Wire notification multiplexer: relay + local WS
            mux = NotificationMultiplexer()
            mux.add_sender(client._send_notification)
            if local_ws_port is not None:
                mux.add_sender(local_server._broadcast_notification)
            message_handler.set_notification_sender(mux.send)

            # Advertise local WS port and token in heartbeat metadata
            if local_ws_port is not None:
                local_ws_token = local_server.token
                client._metadata_supplier = lambda: {
                    "local_ws_port": local_ws_port,
                    "local_ws_token": local_ws_token,
                }

            # Send initial heartbeat so DB has local WS port before relay
            # connection triggers frontend activity
            try:
                await client._send_heartbeat()
            except Exception:
                pass  # Non-fatal — heartbeat loop will retry

            # Start idle monitor for user cloud runners (pool machines skip — no idle timeout)
            is_pool = os.getenv("SKOP_POOL_MACHINE") == "true"
            if is_cloud:
                def idle_shutdown():
                    client.running = False
                    shutdown_event.set()
                    asyncio.create_task(client.disconnect())
                message_handler._shutdown_callback = idle_shutdown
                if not is_pool:
                    message_handler.start_idle_monitor()

                # Reconnect callback for pool machine reassignment
                def force_reconnect():
                    async def _close_ws():
                        if client.ws:
                            await client.ws.close(1000, "Pool reassignment")
                    asyncio.create_task(_close_ws())
                message_handler._reconnect_callback = force_reconnect

            # Start periodic state snapshot loop (every 30s)
            start_snapshot_loop(message_handler, client, output_buffer)

            # Let run_with_reconnect own the relay connection lifecycle
            await client.run_with_reconnect()

        except Exception as e:
            if not shutdown_event.is_set():
                console.print(f"[red]Error: {e}[/red]")
                raise
        finally:
            # Cleanup
            if not shutdown_event.is_set():
                await client.disconnect()
            if local_server is not None:
                await local_server.stop()
            # Final flush of log shipper before exit
            if log_shipper:
                await log_shipper.stop()
            await jupyter_backend.close(preserve_server=True)
            console.print("[green]Runner stopped.[/green]")

    try:
        asyncio.run(run())
    except KeyboardInterrupt:
        pass


@cli.command()
def status():
    """Check runner connection status."""
    from . import __version__
    from .service import get_service_status

    config = RunnerConfig.load()

    if not config.device_token:
        console.print("[red]Not authenticated. Run 'skop-runner login' first.[/red]")
        return

    service = get_service_status()
    service_display = {
        "running": "[green]installed (running)[/green]",
        "stopped": "[yellow]installed (stopped)[/yellow]",
        "not installed": "[dim]not installed[/dim]",
        "not supported": "[dim]not supported on this platform[/dim]",
    }.get(service, service)

    console.print(
        Panel.fit(
            f"Skop Runner v{__version__}\n\n"
            f"Service:    {service_display}\n"
            f"API URL:    {config.api_url}\n"
            f"WS URL:     {config.ws_url}",
            title="Runner Status",
        )
    )


@cli.command()
def install():
    """Install runner as a system service (auto-start on login)."""
    from .service import install_service

    config = RunnerConfig.load()
    if not config.device_token:
        console.print("[red]Not authenticated. Run 'skop-runner login' first.[/red]")
        raise click.Abort()

    try:
        msg = install_service()
        console.print(f"[green]{msg}[/green]")
    except Exception as e:
        console.print(f"[red]Failed to install service: {e}[/red]")
        raise click.Abort()


@cli.command()
def uninstall():
    """Remove runner system service."""
    from .service import uninstall_service

    try:
        msg = uninstall_service()
        console.print(f"[green]{msg}[/green]")
    except Exception as e:
        console.print(f"[red]Failed to uninstall service: {e}[/red]")
        raise click.Abort()


@cli.command()
def logout():
    """Remove authentication."""
    config = RunnerConfig.load()
    config.device_token = None
    config.save()

    console.print("[green]Logged out successfully[/green]")


if __name__ == "__main__":
    cli()
