"""Rust backend bridge — serialize workflow config and spawn Axum binary.

This module handles the `backend="rust"` path for HushApp.serve().
When selected, the Python side serializes graph definitions + endpoint
configs to JSON and spawns a standalone Rust HTTP server (hush-serve)
that executes workflows natively without Python.

Architecture:
    Python (hush-serve)              Rust (hush-serve)
    |- Define graphs via @graph      |- Parse workflow JSON
    |- Register endpoints            |- Load cdylib plugins
    |- Serialize to JSON ----------> |- Build Axum routes
    |- Spawn Rust binary             |- Serve HTTP + SSE + WS
    '- Wait for process              '- Return JSON results

The Rust binary (hush-serve) is a separate crate that:
  - Accepts a config JSON file via --config CLI arg
  - Loads cdylib plugin via --plugin CLI arg
  - Creates Axum routes matching the endpoint definitions
  - Runs entirely without Python — no PyO3 dependency
"""

from __future__ import annotations

import json
import os
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from hush.serve.app import HushApp

# Default location of the hush-serve crate relative to hush-serve package
_HUSH_SERVE_CRATE = Path(__file__).resolve().parents[3] / ".." / "rust" / "hush-serve"


class RustBackendConfig:
    """Serializable configuration for the Rust backend."""

    def __init__(self):
        self.endpoints: List[Dict[str, Any]] = []
        self.host: str = "0.0.0.0"
        self.port: int = 8000
        self.tracers: Optional[Dict[str, Any]] = None
        self.resources: Optional[Dict[str, Any]] = None

    def add_endpoint(self, path: str, graph_config: dict, **kwargs):
        self.endpoints.append({"path": path, "graph": graph_config, **kwargs})

    def to_json(self) -> str:
        data: Dict[str, Any] = {
            "host": self.host,
            "port": self.port,
            "endpoints": self.endpoints,
        }
        if self.tracers:
            data["tracers"] = self.tracers
        if self.resources:
            data["resources"] = self.resources
        return json.dumps(data, default=str)


def serialize_for_rust(app: "HushApp") -> RustBackendConfig:
    """Serialize a HushApp's endpoints into a Rust-compatible config.

    Converts all registered endpoints' graphs via graph.serialize()
    and bundles them with endpoint routing configuration.
    Also extracts tracer configurations for the Rust backend.
    """
    config = RustBackendConfig()

    for ep in app._endpoints:
        graph_config = ep.graph.serialize()
        config.add_endpoint(
            path=ep.config.path,
            graph_config=graph_config,
            stream=ep.config.stream,
            websocket=ep.config.websocket,
        )

    # Extract tracer configs from app-level tracer
    config.tracers = _extract_tracer_configs(app._tracer)

    # Pass through Rust resource configs
    if hasattr(app, "_rust_resources") and app._rust_resources:
        config.resources = app._rust_resources

    return config


def _extract_tracer_configs(tracer) -> Optional[Dict[str, Any]]:
    """Extract Rust-compatible tracer configs from Python Tracer instances."""
    if tracer is None:
        return None

    tracers = tracer if isinstance(tracer, list) else [tracer]
    result: Dict[str, Any] = {}

    for t in tracers:
        cls_name = type(t).__name__

        if cls_name == "LangfuseTracer" and hasattr(t, "_config") and t._config is not None:
            cfg = t._config
            result["langfuse"] = {
                "public_key": cfg.public_key,
                "secret_key": cfg.secret_key,
                "host": cfg.host,
            }
            if hasattr(t, "_stream_trace_limit") and t._stream_trace_limit is not None:
                result["langfuse"]["stream_trace_limit"] = t._stream_trace_limit

        elif cls_name == "HushEyesTracer" and hasattr(t, "_host"):
            result["hush_eyes"] = {
                "host": getattr(t, "_host", "127.0.0.1"),
                "port": getattr(t, "_port", 8420),
            }

    return result if result else None


def find_hush_serve_binary(crate_dir: Optional[Path] = None) -> Path:
    """Find or build the hush-serve binary.

    Search order:
    1. Pre-built binary in crate_dir (monorepo development)
    2. hush-serve on PATH (installed via ``cargo install hush-serve``)
    3. Build from source if crate_dir exists

    Returns:
        Path to the hush-serve binary.

    Raises:
        FileNotFoundError: If hush-serve cannot be found anywhere.
        RuntimeError: If cargo build fails.
    """
    import shutil

    crate = (crate_dir or _HUSH_SERVE_CRATE).resolve()

    if not crate.exists():
        # Crate dir not found — check PATH (e.g. installed via cargo install)
        on_path = shutil.which("hush-serve")
        if on_path:
            return Path(on_path)
        raise FileNotFoundError(
            "hush-serve binary not found. Either:\n"
            "  - Install it: cargo install hush-serve\n"
            "  - Or ensure the hush-serve crate exists in the monorepo"
        )

    # Check for pre-built binary
    if sys.platform == "win32":
        binary_name = "hush-serve.exe"
    else:
        binary_name = "hush-serve"

    # Cargo workspace puts binaries in the workspace root target dir
    workspace_target = crate.parent / "target"
    release_bin = workspace_target / "release" / binary_name
    debug_bin = workspace_target / "debug" / binary_name

    # Prefer release build if it exists
    if release_bin.exists():
        return release_bin
    if debug_bin.exists():
        return debug_bin

    # Build from source (run from workspace root)
    workspace_root = crate.parent
    print(f"Building hush-serve from {workspace_root}...")
    result = subprocess.run(
        ["cargo", "build", "--release", "-p", "hush-serve"],
        cwd=str(workspace_root),
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(f"Failed to build hush-serve:\n{result.stderr}")
    print("hush-serve built successfully.")

    if not release_bin.exists():
        raise RuntimeError(f"Build succeeded but binary not found at {release_bin}")

    return release_bin


def _find_cdylib(rust_ops_path: Path) -> Optional[Path]:
    """Find a built cdylib in a Rust ops crate's target directory."""
    rust_ops_path = rust_ops_path.resolve()
    target_dir = rust_ops_path / "target"

    if sys.platform == "win32":
        ext = ".dll"
    elif sys.platform == "darwin":
        ext = ".dylib"
    else:
        ext = ".so"

    for profile in ("release", "debug"):
        profile_dir = target_dir / profile
        if profile_dir.exists():
            for f in profile_dir.iterdir():
                if f.suffix == ext and f.is_file():
                    return f

    return None


def _build_cdylib(rust_ops_path: Path) -> Path:
    """Build a cdylib from a Rust ops crate and return the path to the built library."""
    rust_ops_path = rust_ops_path.resolve()

    print(f"Building cdylib plugin from {rust_ops_path}...")
    result = subprocess.run(
        ["cargo", "build", "--release"],
        cwd=str(rust_ops_path),
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(f"Failed to build cdylib plugin:\n{result.stderr}")
    print("Plugin built successfully.")

    cdylib = _find_cdylib(rust_ops_path)
    if cdylib is None:
        raise RuntimeError(f"Build succeeded but no cdylib found in {rust_ops_path / 'target'}")
    return cdylib


def serve_rust(
    app: "HushApp",
    host: str,
    port: int,
    crate_dir: Optional[Path] = None,
    rust_ops: Optional[str] = None,
    plugin: Optional[str] = None,
    binary: Optional[str] = None,
):
    """Spawn the Rust backend server.

    Two modes:

    **Custom binary** (recommended for custom ops):
        Pass ``binary=`` pointing to a Rust binary built with ``hush-serve`` as a library.
        The binary already has all ops compiled in — no cdylib/plugin needed.

    **Default binary + plugin** (legacy):
        Uses the generic ``hush-serve`` binary with an optional cdylib plugin.

    Args:
        app: The HushApp instance with registered endpoints.
        host: Bind address.
        port: Bind port.
        crate_dir: Optional path to hush-serve crate (auto-detected if None).
        rust_ops: Optional path to a Rust ops crate (auto-builds cdylib). Legacy.
        plugin: Optional path to a pre-built cdylib plugin (.dll/.so/.dylib). Legacy.
        binary: Optional path to a custom Rust binary (built with hush-serve as library).
    """
    # 1. Serialize config
    config = serialize_for_rust(app)
    config.host = host
    config.port = port
    config_json = config.to_json()

    # 2. Write to temp file
    with tempfile.NamedTemporaryFile(
        mode="w",
        suffix=".json",
        prefix="hush-serve-config-",
        delete=False,
    ) as f:
        f.write(config_json)
        config_path = f.name

    if binary:
        # Custom binary mode: ops are compiled in, no plugin needed
        bin_path = Path(binary).resolve()
        if not bin_path.exists() and sys.platform == "win32" and not bin_path.suffix:
            bin_path = bin_path.with_suffix(".exe")
        if not bin_path.exists():
            raise FileNotFoundError(f"Custom binary not found: {binary}")

        print(f"Starting custom Rust backend at http://{host}:{port}")
        print(f"Binary: {bin_path}")
        print(f"Config: {config_path}")

        cmd = [str(bin_path), "--config", config_path]
    else:
        # Default binary mode: find hush-serve + optional plugin
        plugin_path: Optional[str] = None
        if plugin:
            plugin_path = str(Path(plugin).resolve())
        elif rust_ops:
            ops_path = Path(rust_ops)
            cdylib = _find_cdylib(ops_path)
            if cdylib is None:
                cdylib = _build_cdylib(ops_path)
            plugin_path = str(cdylib)

        bin_path = find_hush_serve_binary(crate_dir)

        print(f"Starting hush-serve (Rust backend) at http://{host}:{port}")
        print(f"Binary: {bin_path}")
        print(f"Config: {config_path}")
        if plugin_path:
            print(f"Plugin: {plugin_path}")

        cmd = [str(bin_path), "--config", config_path]
        if plugin_path:
            cmd.extend(["--plugin", plugin_path])

    # Replace Python process with Rust binary.
    # execvp never returns — Python is gone, Rust owns the process.
    # stdout/stderr are inherited from the parent (bench.py redirects to log files,
    # interactive use goes to terminal).
    os.execvp(cmd[0], cmd)
