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

from __future__ import annotations

import inspect
import os
import subprocess
import tempfile
import textwrap
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.client import VelarClient
from velar.config import Config
from velar.decorators import EndpointSpec, FunctionSpec

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="A10", 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) -> dict[str, Any]:
        """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
        """
        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 {}

        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, Any] = {}

        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)
            if not image_tag:
                console.print(f"  [red]✗ Failed to build image for {spec.name}[/red]")
                continue

            # Step 2: Create deployment (with spinner)
            with Live(
                Text.assemble(("  ⟳ Creating deployment...", "yellow")),
                console=console,
                refresh_per_second=8,
                transient=True,
            ):
                try:
                    deployment = 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),
                    )
                except Exception as e:
                    console.print(f"  [red]✗ Failed: {e}[/red]")
                    continue

            results[spec.name] = deployment

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

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

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

        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.
        func_source = textwrap.dedent(inspect.getsource(spec.func))
        func_name = spec.name

        parts = [
            f"# Auto-generated Velar handler for function '{func_name}'",
            "# " + "-" * 55,
            "import json",
            "import sys",
            "import traceback",
            "from http.server import HTTPServer, BaseHTTPRequestHandler",
            "",
            "# ---------- user function ----------",
            func_source.rstrip(),
            "# -----------------------------------",
            "",
            "",
            "class _VelarHandler(BaseHTTPRequestHandler):",
            "    def do_POST(self):",
            "        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)",
            '            response = json.dumps({"result": result})',
            "            self.send_response(200)",
            "        except Exception:",
            '            response = json.dumps({"error": traceback.format_exc()})',
            "            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:
        """Build Docker image and push to registry. Returns image tag."""
        tag = f"velar/{self.name}/{spec.name}:latest"

        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"

        handler_source = self._generate_handler(spec)

        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)

            with Live(
                Text.assemble(("  ⟳ Building image ", "yellow"), (tag, "cyan"), ("...", "yellow")),
                console=console,
                refresh_per_second=8,
                transient=True,
            ):
                try:
                    result = subprocess.run(
                        ["docker", "build", "-t", tag, "-f", dockerfile_path, tmpdir],
                        capture_output=True,
                        text=True,
                        timeout=300,
                    )
                except FileNotFoundError:
                    console.print("  [red]✗ Docker not found. Install Docker to deploy.[/red]")
                    return None
                except subprocess.TimeoutExpired:
                    console.print("  [red]✗ Docker build timed out (5 min)[/red]")
                    return None

            if result.returncode != 0:
                console.print(f"  [red]✗ Docker build failed:[/red]\n{result.stderr[:400]}")
                return None
            console.print(f"  [green]✓ Image built[/green]  [dim]{tag}[/dim]")

            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 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)
