"""Abstract request handler interface."""

from __future__ import annotations

import hashlib
import multiprocessing as mp
from abc import ABC, abstractmethod
from pathlib import Path
from typing import TYPE_CHECKING, Any

import yaml

from ...context import Event, RequestContext
from ..adapters import validate_adapter_key
from .types import Request, RequestStatus, Response, ResponseAdapterInfo, StreamChunk


class AdapterError(Exception):
    """Raised when adapter resolution fails."""


def _stable_adapter_int_id(adapter: str) -> int:
    """Generate deterministic integer ID from adapter key.

    Uses SHA-256 hash to ensure consistent IDs across process restarts.
    Python's built-in hash() is randomized per-process since Python 3.3.

    Args:
        adapter: The adapter key string.

    Returns:
        Positive 31-bit integer suitable for vLLM's lora_int_id.
    """
    return int(hashlib.sha256(adapter.encode()).hexdigest(), 16) % (2**31)


if TYPE_CHECKING:
    from appinfra.log import Logger

    from ...engines.protocol import InferenceEngineProtocol
    from ..adapters import AdapterManager


class RequestHandler(ABC):
    """
    Abstract base for request execution strategies.

    Implementations define how inference requests are queued, scheduled,
    and executed. The engine loop calls submit() to add requests and
    step() to advance processing.

    For streaming requests, handlers need access to response_q to send
    incremental tokens. Call set_response_queue() before processing.

    Template Method: _process_request() defines the request processing
    algorithm, with subclasses implementing submit() and step() for
    queue management.

    Subclasses must call super().__init__() in their __init__ method.
    """

    def __init__(self) -> None:
        """Initialize base handler state.

        Subclasses must call super().__init__() to ensure proper initialization.
        """
        self._response_q: mp.Queue | None = None
        self._lg: Logger | None = None
        self._lora_base_path: Path | None = None
        self._adapter_manager: AdapterManager | None = None
        self._loaded_adapters: set[str] | None = None

    def set_logger(self, lg: Logger) -> None:
        """Set the logger for request context creation."""
        self._lg = lg

    def set_adapter_manager(self, manager: AdapterManager | None) -> None:
        """Set the adapter manager for validation.

        When set, adapter key must be registered (enabled) in the manager
        for inference to proceed. This enforces the `enabled` field in
        adapter config.yaml files.

        Args:
            manager: AdapterManager instance for validation, or None to
                skip enabled-check (fall back to path-only validation).
        """
        self._adapter_manager = manager

    def get_adapter_manager(self) -> AdapterManager | None:
        """Get the adapter manager for external access (e.g., processors)."""
        return self._adapter_manager

    def set_lora_base_path(self, path: str | None) -> None:
        """Set the base path for LoRA adapter resolution.

        Args:
            path: Base directory for adapter weights. Adapter paths are
                constructed as base_path / adapter_key.
        """
        self._lora_base_path = Path(path).expanduser() if path else None

    def set_response_queue(self, response_q: mp.Queue) -> None:
        """
        Set the response queue for streaming support.

        This allows handlers to put StreamChunk objects directly on the
        queue during generation, enabling token-by-token streaming.

        Args:
            response_q: Queue to send responses/chunks to API layer.
        """
        self._response_q = response_q

    @abstractmethod
    def submit(self, request: Request) -> bool:
        """
        Submit a request for processing.

        Args:
            request: The inference request.

        Returns:
            True if accepted, False if rejected (at capacity).
        """
        pass

    @abstractmethod
    def step(self) -> list[Response]:
        """
        Execute one processing step.

        This is called repeatedly by the engine loop. Each call may:
        - Process one request (Sequential)
        - Process one request from queue (BoundedQueue)
        - Advance all running requests by one token (ContinuousBatching)

        For streaming requests, handlers should put StreamChunk objects
        on self._response_q during processing and only return the final
        Response in the returned list.

        Returns:
            List of completed or rejected responses (may be empty).
        """
        pass

    @property
    @abstractmethod
    def pending_count(self) -> int:
        """Number of requests waiting or in progress."""
        pass

    @property
    @abstractmethod
    def is_saturated(self) -> bool:
        """True if handler cannot accept more requests."""
        pass

    def sequence_stats(self) -> dict[str, int]:
        """Return active sequence statistics.

        Override in subclasses to provide accurate counts.

        Returns:
            Dict with 'active' (number of sequences) and 'total_tokens'.
        """
        return {"active": 0, "total_tokens": 0}

    @property
    @abstractmethod
    def engine(self) -> InferenceEngineProtocol:
        """The inference engine used by this handler."""
        pass

    # -------------------------------------------------------------------------
    # Template Method: Request Processing (shared by all handlers)
    # -------------------------------------------------------------------------

    def _create_context(self, request: Request) -> RequestContext | None:
        """Create RequestContext for a request if logger is available."""
        if self._lg is None:
            return None
        return RequestContext(id=request.id, lg=self._lg)

    def _process_request(self, request: Request) -> Response:
        """Process a single request and return response (Template Method)."""
        ctx = self._create_context(request)
        request.context = ctx
        if ctx:
            ctx.mark(
                Event.REQUESTED, stream=request.stream, max_tokens=request.max_tokens
            )

        if request.stream and self._response_q is not None:
            return self._process_streaming_request(request)
        return self._process_blocking_request(request)

    def _validate_adapter_path(self, adapter: str) -> Path:
        """Validate adapter key and return resolved path.

        Raises:
            AdapterError: If LoRA not configured or adapter key is invalid.
        """
        if not self._lora_base_path:
            raise AdapterError(
                f"adapter '{adapter}' specified but LoRA not configured "
                "(lora.base_path not set)"
            )

        adapter_path = validate_adapter_key(adapter, self._lora_base_path)
        if adapter_path is None:
            if self._lg:
                self._lg.warning(
                    "rejected invalid adapter key", extra={"adapter": adapter}
                )
            raise AdapterError(
                f"invalid adapter '{adapter}': must be a simple name without "
                "path separators or parent references"
            )
        return adapter_path

    def _check_adapter_enabled(self, adapter: str, adapter_path: Path) -> None:
        """Check adapter is enabled for use.

        If AdapterManager is set, uses its cached enabled-adapters list.
        Otherwise falls back to reading config.yaml from disk (useful for
        testing or when manager is not configured).

        Supports versioned adapter resolution: both full keys and base keys
        are accepted when using the adapter manager.

        To hot-reload adapters, use the /v1/adapters/refresh endpoint.

        Raises:
            AdapterError: If adapter is not enabled or config is missing/invalid.
        """
        # Use manager if available (preferred - avoids per-request disk I/O)
        if self._adapter_manager is not None:
            # resolve() handles both full keys and base keys
            if self._adapter_manager.resolve(adapter) is None:
                raise AdapterError(
                    f"adapter '{adapter}' is not enabled. "
                    "Enable it in the adapter's config.yaml and call /v1/adapters/refresh."
                )
            return

        # Fallback: read config.yaml directly (for testing or when manager not set)
        self._check_adapter_config_file(adapter, adapter_path)

    def _check_adapter_config_file(self, adapter: str, adapter_path: Path) -> None:
        """Check adapter config.yaml exists and has enabled: true."""
        config_path = adapter_path / "config.yaml"
        if not config_path.exists():
            raise AdapterError(
                f"adapter '{adapter}' has no config.yaml. "
                "Create config.yaml with 'enabled: true'."
            )

        try:
            with open(config_path, encoding="utf-8") as f:
                config = yaml.safe_load(f) or {}
        except Exception as e:
            raise AdapterError(f"failed to read adapter '{adapter}' config: {e}") from e

        if not isinstance(config, dict):
            raise AdapterError(
                f"adapter '{adapter}' config.yaml must be a mapping, "
                f"got {type(config).__name__}"
            )

        if not config.get("enabled", True):
            raise AdapterError(
                f"adapter '{adapter}' is disabled. "
                "Remove 'enabled: false' from the adapter's config.yaml to enable it."
            )

    def _import_lora_request_class(self, adapter: str) -> type[Any]:
        """Import and return vLLM LoRARequest class, raising AdapterError if unavailable."""
        try:
            from vllm.lora.request import LoRARequest

            return LoRARequest  # type: ignore[no-any-return]
        except ImportError as e:
            if self._lg:
                self._lg.warning(
                    "vLLM LoRA module not available", extra={"adapter": adapter}
                )
            raise AdapterError(
                f"adapter '{adapter}' specified but vLLM LoRA module not available"
            ) from e

    def _log_and_track_adapter(self, adapter: str, adapter_path: Path) -> None:
        """Log adapter usage and track as loaded. Warns on first load."""
        if self._loaded_adapters is None:
            self._loaded_adapters = set()
        is_first_load = adapter not in self._loaded_adapters
        if self._lg:
            # Build extra dict - use "key" when we have full metadata, "adapter" otherwise
            extra: dict[str, Any] = {}
            if self._adapter_manager:
                # Use resolve() to handle both full keys and base keys
                loaded = self._adapter_manager.resolve(adapter)
                if loaded:
                    extra["key"] = loaded.key
                    extra["adapter_name"] = loaded.name
                    extra["md5"] = loaded.md5
                    extra["mtime"] = loaded.mtime
                    extra["peft_type"] = loaded.peft_type or "LORA"
                else:
                    extra["adapter"] = adapter
            else:
                extra["adapter"] = adapter
            if is_first_load:
                extra["path"] = str(adapter_path)
                self._lg.info(
                    "loading PEFT adapter (first use, kernel compilation may take 1-2 min)",
                    extra=extra,
                )
            else:
                self._lg.debug("using PEFT adapter", extra=extra)
        self._loaded_adapters.add(adapter)

    def _resolve_lora_request(
        self, adapter: str | None
    ) -> tuple[Any | None, str | None]:
        """Resolve adapter key to engine request.

        For vLLM engines, returns a LoRARequest object.
        For PEFT engine, returns the adapter path string directly.

        Returns:
            Tuple of (request, fallback_adapter):
            - (LoRARequest or path_str, None) if adapter resolved successfully
            - (None, None) if no adapter requested
            - (None, adapter) if adapter not found (fallback to base model)

        Raises:
            AdapterError: If adapter is empty string or adapter is disabled.
        """
        if adapter is None:
            return None, None
        if not adapter:
            raise AdapterError("adapter cannot be empty")

        adapter_path = self._validate_adapter_path(adapter)
        if not adapter_path.exists():
            if self._lg:
                self._lg.warning(
                    "adapter not found, falling back to base model",
                    extra={"adapter": adapter, "path": str(adapter_path)},
                )
            return None, adapter  # Fallback to base model
        self._check_adapter_enabled(adapter, adapter_path)
        self._log_and_track_adapter(adapter, adapter_path)

        # Try to create vLLM LoRARequest, fall back to path string for PEFT
        try:
            lora_request = self._import_lora_request_class(adapter)(
                lora_name=adapter,
                lora_int_id=_stable_adapter_int_id(adapter),
                lora_path=str(adapter_path),
            )
            return lora_request, None
        except AdapterError:
            # vLLM not available - return path string for PEFT engine
            return str(adapter_path), None

    # Reserved model names that should not be looked up as adapters
    _RESERVED_MODEL_NAMES = frozenset({"auto", "default"})

    def _resolve_effective_adapter(self, request: Request) -> str | None:
        """Resolve effective adapter from request.

        Priority:
        1. Explicit `adapter` field (our extension)
        2. `model` field if it matches a known adapter (OpenAI compatibility)

        Adapter resolution supports versioned keys:
        - Full key (e.g., "my-adapter-a1b2c3d4e5f6") returns exact match
        - Base key (e.g., "my-adapter") resolves to latest version by mtime

        Reserved model names (auto, default) use the base model, not adapters.

        Returns:
            The full (versioned) adapter key to use, or None for base model.
        """
        if request.adapter:
            # Resolve to full key if base key was provided
            if self._adapter_manager:
                resolved = self._adapter_manager.resolve(request.adapter)
                if resolved:
                    return resolved.key
            return request.adapter  # Fall through to validation later

        # Fallback: check if model field matches a known adapter
        # Skip reserved names like "auto" and "default"
        if request.model and self._adapter_manager:
            if request.model not in self._RESERVED_MODEL_NAMES:
                resolved = self._adapter_manager.resolve(request.model)
                if resolved:
                    return resolved.key
        return None

    def _build_engine_params(
        self, request: Request
    ) -> tuple[dict[str, Any], str | None]:
        """Build parameters for engine generate/stream methods from request.

        Returns:
            Tuple of (params_dict, fallback_adapter):
            - params_dict: Parameters for engine generate/stream methods
            - fallback_adapter: Set if adapter was requested but not found
        """
        params: dict[str, Any] = {
            "prompt": request.prompt,
            "max_tokens": request.max_tokens,
            "temperature": request.temperature,
            "top_p": request.top_p,
            "top_k": request.top_k,
            "repetition_penalty": request.repetition_penalty,
            "use_chat_template": request.use_chat_template,
            "stop_sequences": request.stop_sequences,
            "context": request.context,
            "messages": request.messages,
        }
        effective_adapter = self._resolve_effective_adapter(request)
        adapter_result, fallback_adapter = self._resolve_lora_request(effective_adapter)

        if adapter_result is not None:
            # Works for both vLLM LoRARequest and PEFT path string
            params["lora_request"] = adapter_result

        # Tool calling support
        if request.tools:
            params["tools"] = request.tools
        if request.tool_choice is not None:
            params["tool_choice"] = request.tool_choice
        # Structured output support
        if request.response_format is not None:
            params["response_format"] = request.response_format
        return params, fallback_adapter

    def _parse_generate_result(
        self, result: str | dict[str, Any]
    ) -> tuple[
        str, list[dict[str, Any]] | None, dict[str, Any] | None, dict[str, Any] | None
    ]:
        """Parse engine generate result.

        Returns:
            Tuple of (content, tool_calls, usage, adapter_info).
            adapter_info is a dict with requested/actual/fallback/mtime/md5 if present.
        """
        if isinstance(result, dict):
            return (
                result.get("content", ""),
                result.get("tool_calls"),
                result.get("usage"),
                result.get("adapter"),
            )
        return result, None, None, None

    def _build_success_response(
        self,
        request: Request,
        result_text: str,
        tool_calls: list[dict[str, Any]] | None,
        usage: dict[str, Any] | None = None,
        adapter_info: ResponseAdapterInfo | None = None,
    ) -> Response:
        """Build successful response with token counts.

        Uses engine-reported usage when available, falls back to estimation.
        """
        if usage:
            prompt_tokens = usage.get("prompt_tokens", 0)
            completion_tokens = usage.get("completion_tokens", 0)
        else:
            prompt_tokens = self.engine.count_tokens(request.prompt)
            completion_tokens = self.engine.count_tokens(result_text)
        if request.context:
            request.context.mark(
                Event.COMPLETE,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
            )
        return Response(
            id=request.id,
            status=RequestStatus.COMPLETED,
            result=result_text,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            tool_calls=tool_calls,
            adapter=adapter_info,
        )

    def _process_blocking_request(self, request: Request) -> Response:
        """Process request with blocking generation (non-streaming)."""
        try:
            params, fallback_adapter = self._build_engine_params(request)
            result = self.engine.generate(**params)
            if request.context:
                request.context.mark(Event.DECODED)
            result_text, tool_calls, usage, adapter_dict = self._parse_generate_result(
                result
            )
            # Build adapter info from engine response or pre-request fallback
            adapter_info = self._build_adapter_info_from_result(
                adapter_dict, fallback_adapter
            )
            return self._build_success_response(
                request, result_text, tool_calls, usage, adapter_info
            )
        except AdapterError as e:
            if self._lg:
                self._lg.warning(
                    "adapter error", extra={"request_id": request.id, "error": str(e)}
                )
            return Response(id=request.id, status=RequestStatus.FAILED, error=str(e))
        except Exception as e:
            return Response(id=request.id, status=RequestStatus.FAILED, error=str(e))

    def _stream_tokens_to_queue(self, request: Request, stream: Any) -> None:
        """Stream tokens from generator to response queue."""
        assert self._response_q is not None
        for token in stream:
            chunk = StreamChunk(id=request.id, token=token)
            self._response_q.put(chunk)

    def _build_adapter_info_from_result(
        self,
        adapter_dict: dict[str, Any] | None,
        fallback_adapter: str | None,
    ) -> ResponseAdapterInfo | None:
        """Build AdapterInfo from engine result or pre-request fallback.

        Args:
            adapter_dict: Adapter info from engine (has requested/actual/fallback/mtime/md5)
            fallback_adapter: Pre-request fallback (adapter not found before calling engine)
        """
        # Engine returned adapter info (success or mismatch)
        if adapter_dict is not None:
            return ResponseAdapterInfo(
                requested=adapter_dict.get("requested"),
                actual=adapter_dict.get("actual"),
                fallback=adapter_dict.get("fallback", False),
                mtime=adapter_dict.get("mtime"),
                md5=adapter_dict.get("md5"),
            )
        # Pre-request fallback (adapter not available)
        if fallback_adapter is not None:
            return ResponseAdapterInfo(requested=fallback_adapter, fallback=True)
        return None

    def _build_adapter_info(
        self,
        fallback_adapter: str | None,
        stream: Any | None = None,
    ) -> ResponseAdapterInfo | None:
        """Build AdapterInfo from fallback adapter and/or stream mismatch.

        Returns None if no adapter was involved in the request.
        """
        # Check if stream reported adapter info (success or mismatch)
        if stream is not None:
            stream_adapter = getattr(stream, "adapter_info", None)
            if stream_adapter is not None:
                return ResponseAdapterInfo(
                    requested=stream_adapter.get("requested"),
                    actual=stream_adapter.get("actual"),
                    fallback=stream_adapter.get("fallback", False),
                    mtime=stream_adapter.get("mtime"),
                    md5=stream_adapter.get("md5"),
                )
            # Legacy: check for mismatch-only flag
            stream_mismatch = getattr(stream, "adapter_mismatch", False)
            if stream_mismatch:
                requested = getattr(stream, "adapter_requested", None)
                return ResponseAdapterInfo(requested=requested, fallback=True)

        if fallback_adapter is not None:
            return ResponseAdapterInfo(requested=fallback_adapter, fallback=True)

        return None

    def _send_stream_final_chunk(
        self,
        request: Request,
        stream: Any,
        adapter_info: ResponseAdapterInfo | None = None,
    ) -> None:
        """Send final chunk with metadata after streaming completes."""
        assert self._response_q is not None
        final_chunk = StreamChunk(
            id=request.id,
            token="",
            is_final=True,
            finish_reason=stream.finish_reason,
            prompt_tokens=stream.prompt_tokens,
            completion_tokens=stream.completion_tokens,
            tool_calls=getattr(stream, "tool_calls", None),
            adapter=adapter_info,
        )
        self._response_q.put(final_chunk)

    def _finalize_stream(
        self,
        request: Request,
        stream: Any,
        fallback_adapter: str | None = None,
    ) -> Response:
        """Finalize streaming: send final chunk, mark context, return response."""
        ctx = request.context
        if ctx:
            ctx.mark(Event.DECODED)
        # Build adapter info once and reuse
        adapter_info = self._build_adapter_info(fallback_adapter, stream)
        self._send_stream_final_chunk(request, stream, adapter_info)
        if ctx:
            ctx.mark(
                Event.COMPLETE,
                prompt_tokens=stream.prompt_tokens,
                completion_tokens=stream.completion_tokens,
            )
        return Response(
            id=request.id,
            status=RequestStatus.COMPLETED,
            prompt_tokens=stream.prompt_tokens,
            completion_tokens=stream.completion_tokens,
            adapter=adapter_info,
        )

    def _process_streaming_request(self, request: Request) -> Response:
        """Process request with streaming generation."""
        try:
            params, fallback_adapter = self._build_engine_params(request)
            stream = self.engine.generate_stream_sync(**params)
            self._stream_tokens_to_queue(request, stream)
            return self._finalize_stream(request, stream, fallback_adapter)
        except AdapterError as e:
            # Adapter validation errors (invalid key, not enabled, not found)
            if self._lg:
                self._lg.warning(
                    "adapter error", extra={"request_id": request.id, "error": str(e)}
                )
            if self._response_q is not None:
                error_chunk = StreamChunk(
                    id=request.id, token="", is_final=True, finish_reason="error"
                )
                self._response_q.put(error_chunk)
            return Response(id=request.id, status=RequestStatus.FAILED, error=str(e))
        except Exception as e:
            if self._response_q is not None:
                error_chunk = StreamChunk(
                    id=request.id, token="", is_final=True, finish_reason="error"
                )
                self._response_q.put(error_chunk)
            return Response(id=request.id, status=RequestStatus.FAILED, error=str(e))
