"""App - the main entry point for Velar deployments."""

from __future__ import annotations

import hashlib
import inspect
import io
import os
import subprocess
import sys
import tarfile
import tempfile
import textwrap
import time
from typing import Any, Callable

from rich.console import Console
from rich.live import Live
from rich.spinner import Spinner
from rich.table import Table
from rich.text import Text

from velar import analytics as ph
from velar.client import VelarClient
from velar.config import Config
from velar.decorators import EndpointSpec, FunctionSpec
from velar.deployment import Deployment

console = Console()

# The filename written into the container image for the generated handler.
_HANDLER_FILENAME = "_velar_handler.py"
_HANDLER_ENTRY_COMMAND = f"python /app/{_HANDLER_FILENAME}"


class App:
    """A Velar application that groups functions and endpoints.

    Usage:
        import velar

        app = velar.App("my-ml-app")

        image = velar.Image.from_registry("pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime")
        image = image.pip_install("transformers", "accelerate")

        @app.function(gpu="A100", image=image)
        def train(data):
            import torch
            ...

        @app.endpoint(gpu="H100", image=image)
        def predict(request):
            ...

        # Deploy everything
        app.deploy()
    """

    def __init__(self, name: str = "default"):
        self.name = name
        self._functions: dict[str, FunctionSpec] = {}
        self._endpoints: dict[str, EndpointSpec] = {}
        self._client: VelarClient | None = None
        self._local_entrypoint: Callable | None = None

    def function(self, **kwargs: Any):
        """Register a serverless function with this app."""
        from velar.decorators import function as fn_decorator

        def wrapper(func):
            spec = fn_decorator(**kwargs)(func)
            self._functions[spec.name] = spec
            return spec

        return wrapper

    def endpoint(self, **kwargs: Any):
        """Register a persistent endpoint with this app."""
        from velar.decorators import endpoint as ep_decorator

        def wrapper(func):
            spec = ep_decorator(**kwargs)(func)
            self._endpoints[spec.name] = spec
            return spec

        return wrapper

    def local_entrypoint(self):
        """Mark a function as the local entry point for ``velar run``.

        Usage::

            @app.local_entrypoint()
            def main():
                result = my_function.remote(some_data)
                print(result)

        Then run with::

            velar run my_module:app
        """

        def decorator(func: Callable) -> Callable:
            self._local_entrypoint = func
            return func

        return decorator

    @property
    def client(self) -> VelarClient:
        if self._client is None:
            self._client = VelarClient(Config.from_env())
        return self._client

    def deploy(self, *, wait: bool = True, timeout: float = 300) -> dict[str, Deployment]:
        """Deploy all registered functions and endpoints.

        This:
        1. Generates handler wrappers for each function/endpoint
        2. Builds Docker images (including the handler) for each unique image spec
        3. Pushes images to registry
        4. Creates deployments via the Velar API
        5. Stores deployment_id and client refs on each spec so .remote() works
        6. Optionally waits for all deployments to reach 'running' state

        Args:
            wait: If True, block until all deployments are running.
            timeout: Maximum seconds to wait for each deployment.

        Returns:
            A dict mapping function/endpoint names to Deployment objects.
        """
        all_specs: list[FunctionSpec | EndpointSpec] = [
            *self._functions.values(),
            *self._endpoints.values(),
        ]

        if not all_specs:
            console.print("[yellow]No functions or endpoints registered.[/yellow]")
            return {}

        _uid = self.client._config.api_key
        ph.capture(_uid, "sdk.deploy.called", {"app": self.name, "functions": len(self._functions), "endpoints": len(self._endpoints)})
        console.print(f"\n[bold blue]Deploying app '{self.name}'[/bold blue]")
        console.print(f"  Functions : {len(self._functions)}")
        console.print(f"  Endpoints : {len(self._endpoints)}")
        console.print()

        results: dict[str, Deployment] = {}

        for spec in all_specs:
            label = f"[bold]{spec.name}[/bold]  [dim]({spec.deployment_type}, {spec.gpu.name})[/dim]"
            console.print(label)

            # Step 1: Build and push Docker image (with handler baked in)
            image_tag = self._build_and_push(spec)  # remote by default, local if VELAR_LOCAL_BUILD=1
            if not image_tag:
                console.print(f"  [red]✗ Failed to build image for {spec.name}[/red]")
                raise RuntimeError(f"Build failed for '{spec.name}'. Fix the error above and re-run.")

            # Step 2: Create deployment (with spinner)
            with Live(
                Text.assemble(("  ⟳ Creating deployment...", "yellow")),
                console=console,
                refresh_per_second=8,
                transient=True,
            ):
                try:
                    deployment_data = self.client.create_deployment(
                        deployment_type=spec.deployment_type,
                        gpu_type=spec.gpu.name,
                        image_url=image_tag,
                        entry_command=_HANDLER_ENTRY_COMMAND,
                        estimated_seconds=getattr(spec, "timeout", 0),
                        env_vars=getattr(spec, "secrets", None) or {},
                        volume_size_gb=getattr(spec, "volume_size_gb", 0),
                        gpu_count=getattr(spec, "gpu_count", 1),
                    )
                except Exception as e:
                    console.print(f"  [red]✗ Failed: {e}[/red]")
                    continue

            deployment_id = deployment_data["id"]
            deploy = Deployment(
                deployment_id=deployment_id,
                name=spec.name,
                data=deployment_data,
                client=self.client,
            )
            results[spec.name] = deploy

            # Wire up the spec so .remote() can find the deployment.
            spec.deployment_id = deployment_id
            spec.endpoint_url = deployment_data.get("endpoint_url")
            spec._client = self.client

            console.print(f"  [green]✓ Deployment {deployment_id[:8]}[/green]")

            # Register persistent slug if requested (endpoints only)
            if isinstance(spec, EndpointSpec) and spec.slug:
                try:
                    slug_data = self.client.create_endpoint_slug(
                        deployment_id=deployment_id,
                        slug=spec.slug,
                        keep_warm=spec.keep_warm,
                    )
                    invoke_url = slug_data.get("invoke_url", "")
                    spec.invoke_url = invoke_url
                    warm_label = " [keep_warm]" if spec.keep_warm else ""
                    console.print(
                        f"  [cyan]↗ Persistent URL:[/cyan]  [bold]{invoke_url}[/bold]{warm_label}"
                    )
                except Exception as e:
                    console.print(f"  [yellow]⚠ Slug registration failed: {e}[/yellow]")

        if results:
            console.print(f"\n[bold green]✓ Deployed {len(results)} workload(s)[/bold green]")

        # Optionally wait for all deployments to be running
        if wait and results:
            console.print(f"\n[bold blue]Waiting for deployments to start...[/bold blue]")
            for name, deploy in results.items():
                try:
                    deploy.wait(timeout=timeout)
                except Exception as e:
                    console.print(f"  [red]✗ {name}: {e}[/red]")

        return results

    # ------------------------------------------------------------------
    # Handler generation
    # ------------------------------------------------------------------

    @staticmethod
    def _generate_handler(spec: FunctionSpec | EndpointSpec) -> str:
        """Generate a Python HTTP handler script for the given spec.

        The generated script:
        - Embeds the user function source directly (no import needed).
        - Starts a stdlib HTTP server on port 8000.
        - Accepts POST with ``{"args": [...], "kwargs": {...}}``.
        - Returns ``{"result": <return_value>}`` as JSON.
        - Provides a GET health-check at any path.
        """
        # Extract the raw source of the decorated function and dedent it
        # so it sits at the top level of the generated script.
        # Strip any decorator lines (e.g. @app.function(...)) from the source
        raw = inspect.getsource(spec.func)
        lines = raw.splitlines(keepends=True)
        # Drop leading decorator lines (lines starting with @ after stripping whitespace)
        while lines and lines[0].lstrip().startswith("@"):
            lines.pop(0)
        func_source = textwrap.dedent("".join(lines))
        func_name = spec.name

        parts = [
            f"# Auto-generated Velar handler for function '{func_name}'",
            "# " + "-" * 55,
            "import io",
            "import json",
            "import sys",
            "import traceback",
            "from http.server import HTTPServer, BaseHTTPRequestHandler",
            "",
            "# ---------- user function ----------",
            func_source.rstrip(),
            "# -----------------------------------",
            "",
            "",
            "class _VelarHandler(BaseHTTPRequestHandler):",
            "    def do_POST(self):",
            "        _buf = io.StringIO()",
            "        _old_stdout = sys.stdout",
            "        sys.stdout = _buf",
            "        try:",
            '            content_length = int(self.headers.get("Content-Length", 0))',
            "            body = json.loads(self.rfile.read(content_length))",
            '            args = body.get("args", [])',
            '            kwargs = body.get("kwargs", {})',
            f"            result = {func_name}(*args, **kwargs)",
            "            sys.stdout = _old_stdout",
            "            logs = _buf.getvalue().splitlines()",
            '            response = json.dumps({"result": result, "logs": logs})',
            "            self.send_response(200)",
            "        except Exception:",
            "            sys.stdout = _old_stdout",
            "            logs = _buf.getvalue().splitlines()",
            '            response = json.dumps({"error": traceback.format_exc(), "logs": logs})',
            "            self.send_response(500)",
            '        self.send_header("Content-Type", "application/json")',
            "        self.end_headers()",
            "        self.wfile.write(response.encode())",
            "",
            "    def do_GET(self):",
            "        self.send_response(200)",
            '        self.send_header("Content-Type", "application/json")',
            "        self.end_headers()",
            "        self.wfile.write(b'{\"status\":\"ready\"}')",
            "",
            "    def log_message(self, fmt, *args):",
            "        sys.stderr.write(fmt % args + '\\n')",
            "",
            "",
            'if __name__ == "__main__":',
            '    server = HTTPServer(("0.0.0.0", 8000), _VelarHandler)',
            '    print("Velar handler listening on 0.0.0.0:8000", flush=True)',
            "    server.serve_forever()",
            "",
        ]
        return "\n".join(parts)

    # ------------------------------------------------------------------
    # Docker build helpers
    # ------------------------------------------------------------------

    def _build_and_push(self, spec: "FunctionSpec | EndpointSpec") -> "str | None":
        """Dispatch to remote or local build based on VELAR_LOCAL_BUILD env var."""
        if os.environ.get("VELAR_LOCAL_BUILD"):
            return self._local_build_and_push(spec)
        return self._remote_build_and_push(spec)

    def _remote_build_and_push(self, spec: "FunctionSpec | EndpointSpec") -> "str | None":
        """Build image remotely via Velar API — no Docker required.

        Polling loop (3s interval, 35 min timeout):
        - POST /api/v1/builds  → 200 (cache hit) or 201 (new build enqueued)
        - GET  /api/v1/builds/{id} → poll until succeeded or failed
        """
        dockerfile_content = spec.image.to_dockerfile()
        dockerfile_content += f"\nCOPY {_HANDLER_FILENAME} /app/{_HANDLER_FILENAME}\n"
        dockerfile_content += f'CMD ["python", "/app/{_HANDLER_FILENAME}"]\n'
        handler_source = self._generate_handler(spec)

        content_hash = hashlib.sha256(
            (dockerfile_content + handler_source).encode()
        ).hexdigest()[:12]

        try:
            build_data = self.client.create_build(
                app_name=self.name,
                function_name=spec.name,
                content_hash=content_hash,
                dockerfile_content=dockerfile_content,
                handler_source=handler_source,
            )
        except Exception as e:
            console.print(f"  [red]✗ Failed to submit remote build: {e}[/red]")
            return None

        # Cache hit — no polling needed
        _uid = self.client._config.api_key
        if build_data.get("status") == "succeeded":
            tag = build_data["image_tag"]
            ph.capture(_uid, "sdk.build.cache_hit", {"app": self.name, "function": spec.name})
            console.print(f"  [green]✓ Image cached[/green]  [dim]{tag}[/dim]")
            return tag

        build_id = build_data["id"]
        _uid = self.client._config.api_key
        ph.capture(_uid, "sdk.build.remote.started", {"app": self.name, "function": spec.name, "build_id": build_id})
        start = time.monotonic()
        _REMOTE_BUILD_TIMEOUT = 35 * 60  # 35 min — 5 min buffer over BuildKit's 30 min

        from rich.live import Live
        from rich.text import Text

        _poll_error: str | None = None
        with Live(console=console, refresh_per_second=4) as live:
            while True:
                elapsed = int(time.monotonic() - start)
                live.update(
                    Text.assemble(
                        ("  ⟳ Building image remotely... ", "yellow"),
                        (f"{elapsed}s", "dim"),
                    )
                )

                try:
                    build_data = self.client.get_build(build_id)
                except Exception as e:
                    _poll_error = f"  [red]✗ Failed to poll build: {e}[/red]"
                    break

                status = build_data.get("status")
                if status == "succeeded":
                    break
                if status == "failed":
                    error = build_data.get("error") or "unknown error"
                    ph.capture(_uid, "sdk.build.remote.failed", {"app": self.name, "function": spec.name, "error": error[:200]})
                    _poll_error = f"  [red]✗ Remote build failed:[/red]\n{error[:300]}"
                    break
                if elapsed > _REMOTE_BUILD_TIMEOUT:
                    _poll_error = f"  [red]✗ Remote build timed out ({_REMOTE_BUILD_TIMEOUT // 60} min)[/red]"
                    break

                time.sleep(3)

        # Print errors after Live exits to avoid output corruption
        if _poll_error:
            console.print(_poll_error)
            return None

        tag = build_data["image_tag"]
        logs = build_data.get("logs") or ""
        last_logs = "\n".join(logs.splitlines()[-5:])
        if last_logs:
            console.print(f"  [dim]{last_logs}[/dim]")
        ph.capture(_uid, "sdk.build.remote.succeeded", {"app": self.name, "function": spec.name, "duration_s": int(time.monotonic() - start)})
        console.print(f"  [green]✓ Image built remotely[/green]  [dim]{tag}[/dim]")
        return tag

    @staticmethod
    def _check_docker() -> str | None:
        """Return an actionable error message if Docker is not ready, or None if OK."""
        # Is the Docker binary installed?
        result = subprocess.run(["docker", "info"], capture_output=True)
        if result.returncode == 0:
            return None

        stderr = result.stderr.decode("utf-8", errors="replace").lower()
        is_mac = sys.platform == "darwin"
        is_win = sys.platform == "win32"

        if "no such file" in stderr or "not found" in stderr or result.returncode == 127:
            if is_mac:
                return (
                    "Docker is not installed.\n"
                    "  Install it with:  brew install --cask docker\n"
                    "  Or download from: https://docs.docker.com/desktop/mac/"
                )
            elif is_win:
                return (
                    "Docker is not installed.\n"
                    "  Download Docker Desktop from: https://docs.docker.com/desktop/windows/"
                )
            else:
                return (
                    "Docker is not installed.\n"
                    "  Install it with:  sudo apt install docker.io  (Ubuntu/Debian)\n"
                    "  Or see:           https://docs.docker.com/engine/install/"
                )

        # Binary exists but daemon is not running
        if is_mac or is_win:
            return (
                "Docker is not running.\n"
                "  Open the Docker Desktop app and wait for it to start, then try again."
            )
        else:
            return (
                "Docker daemon is not running.\n"
                "  Start it with:  sudo systemctl start docker\n"
                "  Then try again."
            )

    def _local_build_and_push(self, spec: "FunctionSpec | EndpointSpec") -> "str | None":
        """Build Docker image locally and push to the Velar registry.

        Activated by VELAR_LOCAL_BUILD=1. Requires Docker installed locally.

        Flow:
        1. Generate handler source and write Dockerfile to a temp dir.
        2. ``docker build`` the image.
        3. ``docker login`` to the Velar registry using the API key.
        4. ``docker tag`` + ``docker push`` to ``registry.velar.run/…``.
        """
        docker_error = self._check_docker()
        if docker_error:
            console.print(f"  [red]✗ {docker_error}[/red]")
            return None

        config = self.client._config
        registry = config.registry_url  # e.g. "registry.velar.run"

        dockerfile_content = spec.image.to_dockerfile()

        # Append a COPY instruction for the generated handler so it lands
        # at /app/_velar_handler.py inside the container.
        dockerfile_content += f"\nCOPY {_HANDLER_FILENAME} /app/{_HANDLER_FILENAME}\n"
        dockerfile_content += f'CMD ["python", "/app/{_HANDLER_FILENAME}"]\n'

        handler_source = self._generate_handler(spec)

        # Content-addressed tag: hash of the full Dockerfile + handler source.
        # Identical code + identical deps → same tag → safe to skip rebuild.
        content_hash = hashlib.sha256(
            (dockerfile_content + handler_source).encode()
        ).hexdigest()[:12]
        tag = f"{registry}/{self.name}/{spec.name}:{content_hash}"

        with tempfile.TemporaryDirectory() as tmpdir:
            # Write Dockerfile
            dockerfile_path = os.path.join(tmpdir, "Dockerfile")
            with open(dockerfile_path, "w") as f:
                f.write(dockerfile_content)

            # Write the handler script next to the Dockerfile so COPY works.
            handler_path = os.path.join(tmpdir, _HANDLER_FILENAME)
            with open(handler_path, "w") as f:
                f.write(handler_source)

            # ----- Check local cache -----
            # If an image with this exact content hash already exists locally,
            # skip the (slow) docker build entirely.
            inspect_result = subprocess.run(
                ["docker", "image", "inspect", tag],
                capture_output=True,
                text=True,
            )
            image_cached = inspect_result.returncode == 0

            if image_cached:
                console.print(f"  [green]✓ Image cached[/green]  [dim]{tag}[/dim]")
            else:
                # Build context as in-memory tar so Docker daemon doesn't need
                # filesystem access to the temp dir (fixes BuildKit / remote
                # daemon / macOS symlink issues where the host /tmp path is not
                # visible to the daemon).
                context_buf = io.BytesIO()
                with tarfile.open(fileobj=context_buf, mode="w", format=tarfile.GNU_FORMAT) as tar:
                    tar.add(dockerfile_path, arcname="Dockerfile")
                    tar.add(handler_path, arcname=_HANDLER_FILENAME)
                tar_bytes = context_buf.getvalue()

                # ----- Build -----
                with Live(
                    Text.assemble(
                        ("  ⟳ Building image ", "yellow"),
                        (tag, "cyan"),
                        ("... (large ML images may take 10–15 min on first build)", "yellow"),
                    ),
                    console=console,
                    refresh_per_second=8,
                    transient=True,
                ):
                    try:
                        build_result = subprocess.run(
                            ["docker", "build", "--platform", "linux/amd64", "-t", tag, "-"],
                            input=tar_bytes,
                            capture_output=True,
                            timeout=1800,  # 30 min — large ML base images can be slow
                        )
                    except FileNotFoundError:
                        console.print("  [red]✗ Docker not found.[/red]")
                        return None
                    except subprocess.TimeoutExpired:
                        console.print(
                            "  [red]✗ Docker build timed out (30 min).[/red]\n"
                            "  This usually means the base image is very large or your internet is slow.\n"
                            "  Try again — layers are cached so subsequent builds will be faster."
                        )
                        return None

                if build_result.returncode != 0:
                    stderr = build_result.stderr.decode("utf-8", errors="replace")
                    console.print(f"  [red]✗ Docker build failed:[/red]\n{stderr[:400]}")
                    return None
                console.print(f"  [green]✓ Image built[/green]  [dim]{tag}[/dim]")

            # ----- Obtain registry token and Docker login -----
            registry_token = self._get_registry_token(
                config, scope=f"repository:{self.name}/{spec.name}:push,pull"
            )
            if not registry_token:
                return None

            login_result = subprocess.run(
                ["docker", "login", registry, "-u", "v0", "--password-stdin"],
                input=registry_token,
                capture_output=True,
                text=True,
                timeout=30,
            )
            if login_result.returncode != 0:
                console.print(
                    f"  [red]✗ Registry login failed:[/red] {login_result.stderr[:200]}"
                )
                return None

            # ----- Push -----
            with Live(
                Text.assemble(("  ⟳ Pushing image...", "yellow")),
                console=console,
                refresh_per_second=8,
                transient=True,
            ):
                try:
                    result = subprocess.run(
                        ["docker", "push", tag],
                        capture_output=True,
                        text=True,
                        timeout=300,
                    )
                except subprocess.TimeoutExpired:
                    console.print("  [red]✗ Docker push timed out (5 min)[/red]")
                    return None

            if result.returncode != 0:
                console.print(f"  [red]✗ Docker push failed:[/red] {result.stderr[:200]}")
                return None
            console.print("  [green]✓ Image pushed[/green]")

        return tag

    def _get_registry_token(self, config: Config, scope: str = "") -> str | None:
        """Fetch a short-lived JWT from the Velar API for Docker registry auth.

        The backend validates the API key and returns an ES256-signed JWT
        that the Cloudflare registry Worker trusts.
        """
        import base64
        import httpx

        basic_creds = base64.b64encode(f"velar:{config.api_key}".encode()).decode()
        try:
            resp = httpx.get(
                f"{config.api_url}/api/v1/registry/token",
                params={
                    "service": config.registry_url,
                    "scope": scope,
                },
                headers={"Authorization": f"Basic {basic_creds}"},
                timeout=15,
            )
            resp.raise_for_status()
            token = resp.json().get("token")
            if not token:
                console.print("  [red]✗ No token in registry response[/red]")
                return None
            return token
        except Exception as e:
            console.print(f"  [red]✗ Failed to get registry token: {e}[/red]")
            return None

    def status(self) -> None:
        """Print status of all deployments."""
        deployments = self.client.list_deployments()

        table = Table(title=f"Deployments - {self.name}")
        table.add_column("ID", style="cyan")
        table.add_column("Type")
        table.add_column("GPU")
        table.add_column("Status")
        table.add_column("Cost")

        for d in deployments:
            status_color = {
                "running": "green",
                "pending": "yellow",
                "provisioning": "yellow",
                "completed": "blue",
                "failed": "red",
                "cancelled": "dim",
            }.get(d["status"], "white")

            actual = d.get("actual_cost")
            cost_str = f"${actual:.4f}" if actual else f"~${d.get('cost_reserved', 0):.4f}"

            table.add_row(
                d["id"][:8],
                d["type"],
                d["gpu_type"],
                f"[{status_color}]{d['status']}[/{status_color}]",
                cost_str,
            )

        console.print(table)
