"""Runner dispatch — registry-driven.

The route layer calls runner_for(model_key) and gets back something
implementing the Runner protocol. Routes don't know which class they
got, only that they can call .run() / .stream() on it.

Resolution:
    1. Look up model_key in MODEL_REGISTRY -> get (framework, task)
    2. Look up (framework, task) in _RUNNERS -> get a runner class
    3. Instantiate (cached per model_key) and return

Adding a new model:
    Add a row to MODEL_REGISTRY. Done — if its (framework, task) pair
    already has a runner class registered.

Adding a new task family:
    1. Define new TaskRequest / TaskResult types
    2. Write a runner class implementing the Runner protocol
    3. Add a row to _RUNNERS for the (framework, task) keys it handles
    4. Add a row to MODEL_REGISTRY for the model
    No route changes needed.

Why a per-process cache instead of recreating runners every call:
    Loading a 14B model takes seconds; doing it on every request is
    obviously wrong. Per-key caching ensures one model = one loaded
    instance per worker process. Inner singletons (REGISTRY for DeepCoder,
    get_llama_runner for llama.cpp) handle further deduplication.
"""
from __future__ import annotations

import logging
import threading
from typing import Dict, Tuple, Type

from .imports import Runner
from .generate import DeepCoderChatRunner
from .vision import VisionRunner
from .llama import LlamaCppChatRunner
# Adjust this import to wherever your model registry actually lives.
# Used only to look up (framework, task) for a given model_key.
from ..imports import MODEL_REGISTRY  # type: ignore

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Dispatch table — single source of truth for (framework, task) -> Runner.
#
# Currently handles chat-family models (transformers causal LMs + llama.cpp
# GGUFs). Other task families plug in here as they're implemented:
#   ("transformers", "summarization")      -> SummarizeRunner
#   ("transformers", "speech-recognition") -> WhisperRunner
#   ("transformers", "vision-language")    -> VisionRunner
#   ("transformers", "embeddings")         -> KeywordRunner / EmbeddingRunner
# ---------------------------------------------------------------------------

_RUNNERS: Dict[Tuple[str, str], Type[Runner]] = {
    ("transformers", "code-generation"): DeepCoderChatRunner,
    ("transformers", "text-generation"): DeepCoderChatRunner,
    ("llama_cpp",    "code-generation"): LlamaCppChatRunner,
    ("llama_cpp",    "text-generation"): LlamaCppChatRunner,
    ("transformers", "vision-language"): VisionRunner

}


# ---------------------------------------------------------------------------
# Per-process instance cache
# ---------------------------------------------------------------------------

_INSTANCES: Dict[str, Runner] = {}
_INSTANCES_LOCK = threading.Lock()


def runner_for(model_key: str) -> Runner:
    """Return the cached Runner for `model_key`, building it on first call.

    Raises KeyError if the model isn't in MODEL_REGISTRY, or if no runner
    class is registered for its (framework, task) pair. Both errors carry
    enough info for the caller to surface a useful 4xx.
    """
    # Fast path: already built.
    cached = _INSTANCES.get(model_key)
    if cached is not None:
        return cached

    with _INSTANCES_LOCK:
        # Re-check after acquiring the lock — another thread may have
        # built it while we were waiting.
        cached = _INSTANCES.get(model_key)
        if cached is not None:
            return cached

        cfg = MODEL_REGISTRY.get(model_key)
        if cfg is None:
            raise KeyError(
                f"Unknown model_key={model_key!r}; "
                f"known: {sorted(MODEL_REGISTRY.keys())}"
            )

        key = (cfg.framework, cfg.task)
        cls = _RUNNERS.get(key)
        if cls is None:
            raise KeyError(
                f"No runner registered for {model_key!r} "
                f"(framework={cfg.framework!r}, task={cfg.task!r}); "
                f"known runner keys: {sorted(_RUNNERS.keys())}"
            )

        logger.info(
            "instantiating runner: model=%s class=%s framework=%s task=%s",
            model_key, cls.__name__, cfg.framework, cfg.task,
        )
        instance = cls(model_key)
        _INSTANCES[model_key] = instance
        return instance


# ---------------------------------------------------------------------------
# Inspection / lifecycle helpers — useful for tests, ops, and debugging
# ---------------------------------------------------------------------------

def loaded_model_keys() -> list[str]:
    """Which model_keys currently have a runner instantiated."""
    with _INSTANCES_LOCK:
        return sorted(_INSTANCES.keys())


def evict(model_key: str) -> bool:
    """Drop a runner from the cache. Returns True if something was evicted.

    The underlying model may still be loaded if the inner singleton
    (REGISTRY for DeepCoder, _LLAMA_INSTANCES for llama.cpp) holds it.
    Eviction here only releases the runner wrapper.
    """
    with _INSTANCES_LOCK:
        return _INSTANCES.pop(model_key, None) is not None


def clear() -> None:
    """Drop all cached runners. Tests use this; production probably shouldn't."""
    with _INSTANCES_LOCK:
        _INSTANCES.clear()


def supported_task_keys() -> list[Tuple[str, str]]:
    """List the (framework, task) pairs the dispatch table currently handles."""
    return sorted(_RUNNERS.keys())
