"""rc serve — run as a provider node, polling for and executing tasks."""

from __future__ import annotations

import signal
import threading
import time

import typer
from rich.console import Console

console = Console()

_shutdown_requested = False
DEFAULT_TASK_LEASE_SECONDS = 300
DEFAULT_TASK_KEEPALIVE_INTERVAL = 30


def _handle_signal(signum: int, frame: object) -> None:
    global _shutdown_requested
    _shutdown_requested = True
    console.print("\n[yellow]Shutdown requested, finishing current task...[/yellow]")


def run(
    model: str = typer.Option(None, "--model", "-m", help="Override default model for inference"),
    poll_interval: int = typer.Option(5, "--poll-interval", "-p", help="Seconds between task polls"),
    max_tokens: int = typer.Option(4096, "--max-tokens", help="Maximum tokens per generation"),
    temperature: float = typer.Option(0.6, "--temperature", "-t", help="Sampling temperature"),
    threads: int | None = typer.Option(
        None,
        "--threads",
        help="Inference threads per provider; defaults to auto-splitting local CPU capacity",
    ),
    ctx_size: int = typer.Option(2048, "--ctx-size", help="llama.cpp context window"),
    gpu_layers: int | None = typer.Option(
        None,
        "--gpu-layers",
        help="GPU layers to offload; defaults to RC_N_GPU_LAYERS or llama.cpp defaults",
    ),
    lease_seconds: int = typer.Option(
        DEFAULT_TASK_LEASE_SECONDS,
        "--lease-seconds",
        help="Task lease duration in seconds; renewed while inference is running",
    ),
    keepalive_interval: int = typer.Option(
        DEFAULT_TASK_KEEPALIVE_INTERVAL,
        "--keepalive-interval",
        help="Seconds between in-task heartbeats and lease renewals",
    ),
) -> None:
    """Start this node as an inference provider, polling for tasks and earning credits."""
    from rc.auth import node_auth_payload
    from rc.config import load_config, save_config
    from rc.credits import format_credit_subunits
    from rc.inference.catalog import get_model_spec
    from rc.inference.flops import format_flops_as_tflops
    from rc.inference.provision import detect_hardware, ensure_model_downloaded
    from rc.supabase_client import get_client

    cfg = load_config()
    if cfg.role == "requester":
        console.print("[red]This node is configured as requester-only. Cannot serve.[/red]")
        raise typer.Exit(code=1)

    model_name = model or cfg.model
    client = get_client()
    if not cfg.node_secret:
        cfg = load_config()
        model_name = model or cfg.model

    signal.signal(signal.SIGINT, _handle_signal)
    signal.signal(signal.SIGTERM, _handle_signal)

    console.print("[bold green]Research Capacitor — Provider Mode[/bold green]")
    console.print(f"  Node:  {cfg.node_name} ({cfg.node_id[:8]}...)")
    console.print(f"  Model: {model_name}")
    console.print(f"  Poll:  every {poll_interval}s")
    console.print(f"  Lease: {lease_seconds}s")
    console.print()

    profile = detect_hardware()
    config_changed = False
    if cfg.accelerator != profile.accelerator:
        cfg.accelerator = profile.accelerator
        config_changed = True
    if cfg.gpu_count != profile.gpu_count:
        cfg.gpu_count = profile.gpu_count
        config_changed = True
    if cfg.gpu_vram_gb != profile.total_vram_gb:
        cfg.gpu_vram_gb = profile.total_vram_gb
        config_changed = True
    if cfg.total_memory_gb != profile.total_memory_gb:
        cfg.total_memory_gb = profile.total_memory_gb
        config_changed = True
    if cfg.gpu_name != profile.gpu_name:
        cfg.gpu_name = profile.gpu_name
        config_changed = True
    if config_changed:
        save_config(cfg)
    console.print(f"[dim]Detected accelerator: {profile.accelerator}[/dim]")
    model_spec = get_model_spec(model_name)
    if model_spec is None:
        console.print(f"[red]Unsupported model '{model_name}'.[/red]")
        raise typer.Exit(code=1)

    try:
        local_model_path = ensure_model_downloaded(model_name, cfg.model_path)
    except Exception as exc:
        console.print(f"[red]Failed to prepare model '{model_name}': {exc}[/red]")
        raise typer.Exit(code=1)

    console.print(f"[dim]Loading {model_spec.runtime} engine...[/dim]")
    from rc.inference.engine import InferenceEngine

    engine = InferenceEngine(
        model=str(local_model_path),
        runtime=model_spec.runtime,
        model_params=model_spec.param_count,
        temperature=temperature,
        max_tokens=max_tokens,
        n_ctx=ctx_size,
        n_threads=threads,
        n_gpu_layers=gpu_layers,
    )
    console.print("[green]Engine loaded.[/green]\n")

    _send_heartbeat(client, cfg, online=True)
    tasks_completed = 0
    total_credits = 0

    try:
        while not _shutdown_requested:
            _send_heartbeat(client, cfg, online=True)

            task_id = _claim_task(client, cfg, lease_seconds)
            if task_id is None:
                _print_idle(tasks_completed, total_credits, poll_interval)
                _sleep_interruptible(poll_interval)
                continue

            task = _fetch_task(client, cfg, task_id)
            if task is None:
                continue

            console.print(f"\n[cyan]Claimed task {task_id[:8]}...[/cyan]")
            console.print(f"  Prompt: {task['prompt'][:80]}...")

            if not _mark_task_running(client, cfg, task_id, lease_seconds):
                console.print(
                    f"[yellow]Task {task_id[:8]}... could not transition to running. Skipping.[/yellow]"
                )
                continue

            task_keepalive = _TaskKeepalive(
                cfg=cfg,
                task_id=task_id,
                lease_seconds=lease_seconds,
                interval_seconds=_normalize_keepalive_interval(
                    keepalive_interval,
                    lease_seconds,
                ),
            )
            task_keepalive.start()

            try:
                task_type = task.get("type", "lean_proof_gen")
                result = engine.generate_lean(
                    task["prompt"],
                    task_type=task_type,
                    context=_build_generation_context(task),
                )
                task_keepalive.stop()
                billed_credits = min(result.credit_subunits, int(task.get("credit_budget", 0)))
                released_credits = max(int(task.get("credit_budget", 0)) - billed_credits, 0)

                client.rpc(
                    "complete_task",
                    {
                        "p_task_id": task_id,
                        "p_lean_code": result.lean_code,
                        "p_raw_output": result.raw_output,
                        "p_mflops": result.mflops,
                        "p_prompt_tokens": result.prompt_tokens,
                        "p_generated_tokens": result.generated_tokens,
                        "p_model_params": result.model_params,
                        "p_prefill_eval_ms": result.prefill_eval_ms,
                        "p_decode_eval_ms": result.decode_eval_ms,
                        "p_total_eval_ms": result.total_eval_ms,
                        **node_auth_payload(cfg),
                    },
                ).execute()

                tasks_completed += 1
                total_credits += billed_credits
                console.print(
                    f"[green]Task {task_id[:8]}... completed![/green] "
                    f"(actual {format_flops_as_tflops(result.flops)} TFLOPs, "
                    f"billed {format_credit_subunits(billed_credits)} credits, "
                    f"{result.generated_tokens} tokens)"
                )
                if released_credits > 0:
                    console.print(
                        f"[dim]Released requester reserve: {format_credit_subunits(released_credits)} credits[/dim]"
                    )
            except Exception as exc:
                task_keepalive.stop()
                console.print(f"[red]Task {task_id[:8]}... failed: {exc}[/red]")
                client.rpc(
                    "fail_task",
                    {
                        "p_task_id": task_id,
                        "p_error": str(exc),
                        **node_auth_payload(cfg),
                    },
                ).execute()

    finally:
        _send_heartbeat(client, cfg, online=False)
        console.print()
        console.print("[bold]Session summary:[/bold]")
        console.print(f"  Tasks completed: {tasks_completed}")
        console.print(f"  Credits earned:  {format_credit_subunits(total_credits)}")
        console.print("[dim]Node status set to offline.[/dim]")


def _send_heartbeat(client: object, cfg: object, online: bool = True) -> None:
    """Update node status and heartbeat timestamp."""
    from rc.auth import node_auth_payload

    try:
        client.rpc(
            "update_node_heartbeat",
            {
                **node_auth_payload(cfg),
                "p_online": online,
                "p_gpu_info": cfg.gpu_info_payload(),
            },
        ).execute()
    except Exception:
        pass


def _claim_task(client: object, cfg: object, lease_seconds: int) -> str | None:
    """Try to atomically claim a pending task via the claim_task RPC."""
    from rc.auth import node_auth_payload

    try:
        params = node_auth_payload(cfg)
        params["p_lease_seconds"] = lease_seconds
        resp = client.rpc("claim_task", params).execute()
        return resp.data if resp.data else None
    except Exception:
        return None


def _mark_task_running(
    client: object,
    cfg: object,
    task_id: str,
    lease_seconds: int,
) -> bool:
    from rc.auth import node_auth_payload

    try:
        resp = client.rpc(
            "mark_task_running",
            {
                **node_auth_payload(cfg),
                "p_task_id": task_id,
                "p_lease_seconds": lease_seconds,
            },
        ).execute()
        return bool(resp.data)
    except Exception:
        return False


def _renew_task_lease(
    client: object,
    cfg: object,
    task_id: str,
    lease_seconds: int,
) -> bool:
    from rc.auth import node_auth_payload

    try:
        resp = client.rpc(
            "renew_task_lease",
            {
                **node_auth_payload(cfg),
                "p_task_id": task_id,
                "p_lease_seconds": lease_seconds,
            },
        ).execute()
        return bool(resp.data)
    except Exception:
        return False


def _fetch_task(client: object, cfg: object, task_id: str) -> dict | None:
    """Fetch full task details by ID."""
    from rc.auth import node_auth_payload

    try:
        resp = client.rpc(
            "get_task",
            {**node_auth_payload(cfg), "p_task_id": task_id},
        ).execute()
        return resp.data
    except Exception:
        return None


def _print_idle(tasks_completed: int, total_credits: int, poll_interval: int) -> None:
    from rc.credits import format_credit_subunits

    console.print(
        f"[dim]Idle — {tasks_completed} tasks done, "
        f"{format_credit_subunits(total_credits)} credits earned. "
        f"Polling again in {poll_interval}s...[/dim]",
        end="\r",
    )


def _sleep_interruptible(seconds: int) -> None:
    """Sleep in small increments so we can respond to SIGINT quickly."""
    for _ in range(seconds * 4):
        if _shutdown_requested:
            break
        time.sleep(0.25)


def _normalize_keepalive_interval(interval_seconds: int, lease_seconds: int) -> int:
    maximum_interval = max(5, lease_seconds // 3)
    return min(max(interval_seconds, 5), maximum_interval)


def _build_generation_context(task: dict) -> dict[str, object]:
    """Merge task, episode, and attempt context into one prompt payload."""
    task_metadata = dict(task.get("task_metadata") or {})
    episode = task.get("episode") or {}
    attempt = task.get("proof_attempt") or {}
    if episode.get("domain"):
        task_metadata.setdefault("domain", episode.get("domain"))
    if episode.get("conjecture_text"):
        task_metadata.setdefault("conjecture_text", episode.get("conjecture_text"))
    if episode.get("formal_statement"):
        task_metadata.setdefault("formal_statement", episode.get("formal_statement"))
    if episode.get("strategy_tag"):
        task_metadata.setdefault("strategy_tag", episode.get("strategy_tag"))
        task_metadata.setdefault("conjecture_family", episode.get("strategy_tag"))
    if attempt.get("generated_lean_code"):
        task_metadata.setdefault("previous_code", attempt.get("generated_lean_code"))
    if attempt.get("lean_stderr"):
        task_metadata.setdefault("lean_stderr", attempt.get("lean_stderr"))
    task_metadata.setdefault("task_metadata", task.get("task_metadata") or {})
    return task_metadata


class _TaskKeepalive:
    """Keep provider liveness and task leases fresh during long inference runs."""

    def __init__(
        self,
        cfg: object,
        task_id: str,
        lease_seconds: int,
        interval_seconds: int,
    ) -> None:
        self._cfg = cfg
        self._task_id = task_id
        self._lease_seconds = lease_seconds
        self._interval_seconds = interval_seconds
        self._lease_renewal_enabled = True
        self._stop_event = threading.Event()
        self._thread = threading.Thread(
            target=self._run,
            name=f"rc-task-keepalive-{task_id[:8]}",
            daemon=True,
        )

    def start(self) -> None:
        self._thread.start()

    def stop(self) -> None:
        self._stop_event.set()
        self._thread.join(timeout=self._interval_seconds + 1)

    def _run(self) -> None:
        from supabase import create_client

        client = create_client(self._cfg.supabase_url, self._cfg.supabase_key)
        while not self._stop_event.wait(self._interval_seconds):
            _send_heartbeat(client, self._cfg, online=True)
            if not self._lease_renewal_enabled:
                continue
            if _renew_task_lease(client, self._cfg, self._task_id, self._lease_seconds):
                continue
            self._lease_renewal_enabled = False
            console.print(
                f"[yellow]Task {self._task_id[:8]}... heartbeat is still active, but lease renewal is unavailable until the SQL update is applied.[/yellow]"
            )
