"""Ollama native client for the shared inference gateway.

Implements generation via the Ollama /api/chat endpoint with:
- Shared StructuredGenerationRequest/InferenceResult types
- Native JSON schema formatting when supported (format field)
- Prompt-only fallback with explicit reporting
- Stall/loop detection as Ollama-specific policy
- Configurable max output tokens (num_predict) and context window (num_ctx)
- Error mapping to shared InferenceErrorCategory

Requirements: 2.1, 2.6
"""
from __future__ import annotations

import json
import logging
import time
from dataclasses import dataclass

import httpx

from services.shared.inference.errors import InferenceError, InferenceErrorCategory
from services.shared.inference.models import (
    InferenceResult,
    InferenceTarget,
    StructuredGenerationRequest,
    TokenUsage,
)

# Re-export for backward compatibility
__all__ = ["OllamaNativeClient", "StallPolicy"]

logger = logging.getLogger("inference.ollama_native")


@dataclass
class StallPolicy:
    """Ollama-specific stall/loop detection configuration.

    Monitors streaming responses for repetitive output patterns.
    This policy is intentionally Ollama-specific and does not leak
    into the generic InferenceResult interface.
    """

    enabled: bool = True
    check_interval_seconds: float = 5.0
    max_unchanged_intervals: int = 6
    loop_window: int = 64
    loop_threshold: float = 0.5


def _map_http_status_to_category(status_code: int) -> InferenceErrorCategory:
    """Map Ollama HTTP status codes to normalized error categories."""
    if status_code == 401:
        return InferenceErrorCategory.AUTH_FAILED
    if status_code == 403:
        return InferenceErrorCategory.FORBIDDEN
    if status_code == 404:
        return InferenceErrorCategory.MODEL_NOT_FOUND
    if status_code == 429:
        return InferenceErrorCategory.RATE_LIMITED
    if status_code in (400, 422):
        return InferenceErrorCategory.BAD_REQUEST
    if status_code >= 500:
        return InferenceErrorCategory.SERVER_ERROR
    return InferenceErrorCategory.UNKNOWN


def _detect_loop(content: str, window: int, threshold: float) -> bool:
    """Detect repetitive output by checking if the tail repeats earlier content.

    Looks at the last `window` characters and checks if a significant
    portion of the tail matches an earlier substring, indicating the model
    is stuck in a generation loop.
    """
    if len(content) < window * 2:
        return False

    tail = content[-window:]
    body = content[:-window]

    # Check if the tail appears verbatim in the preceding content
    if tail in body:
        return True

    # Check character-level repetition ratio in the tail
    if not tail:
        return False
    unique_chars = len(set(tail))
    ratio = unique_chars / len(tail)
    return ratio < threshold


class OllamaNativeClient:
    """Async client for Ollama /api/chat using shared inference types.

    Translates StructuredGenerationRequest into Ollama's native API format
    and returns InferenceResult with proper metadata and error classification.

    Stall/loop detection is an Ollama-specific policy that monitors streaming
    responses for repetitive patterns and aborts generation when detected.
    This does not affect the InferenceResult interface — stall detection
    raises an InferenceError with category STALL_DETECTED.
    """

    def __init__(
        self,
        target: InferenceTarget,
        *,
        http_client: httpx.AsyncClient | None = None,
        stall_policy: StallPolicy | None = None,
    ) -> None:
        if target.protocol != "ollama_native":
            raise ValueError(
                f"OllamaNativeClient requires protocol='ollama_native', got '{target.protocol}'"
            )
        self._target = target
        self._stall_policy = stall_policy or StallPolicy()
        self._owns_client = http_client is None
        self._http = http_client or httpx.AsyncClient(
            timeout=httpx.Timeout(300.0, read=300.0),
        )

    async def close(self) -> None:
        """Close the underlying HTTP client if we own it."""
        if self._owns_client:
            await self._http.aclose()

    async def generate(self, request: StructuredGenerationRequest) -> InferenceResult:
        """Send a structured generation request to Ollama and return the result.

        Builds the Ollama-native payload, streams the response, applies
        stall detection, and maps results to InferenceResult.
        """
        start = time.monotonic()

        payload = self._build_payload(request)
        url = f"{self._target.base_url}/api/chat"

        logger.info(
            "Ollama POST %s model=%s messages=%d max_tokens=%d",
            url,
            self._target.model,
            len(request.messages),
            request.max_output_tokens,
        )

        try:
            content, metadata = await self._stream_response(url, payload, request.timeout_seconds)
        except InferenceError:
            raise
        except httpx.TimeoutException as exc:
            raise InferenceError(
                InferenceErrorCategory.TIMEOUT,
                f"Request timed out after {request.timeout_seconds}s",
                provider_detail=str(exc),
            ) from exc
        except httpx.ConnectError as exc:
            raise InferenceError(
                InferenceErrorCategory.CONNECTION_REFUSED,
                "Connection refused",
                provider_detail=str(exc),
            ) from exc
        except httpx.HTTPError as exc:
            raise InferenceError(
                InferenceErrorCategory.CONNECTION_ERROR,
                "HTTP connection error",
                provider_detail=str(exc),
            ) from exc

        latency_ms = int((time.monotonic() - start) * 1000)

        # Determine structured mode
        structured_mode = self._determine_structured_mode(request)

        # Attempt JSON parsing if we expect structured output
        parsed = None
        if request.json_schema and content:
            try:
                parsed = json.loads(content)
            except (json.JSONDecodeError, ValueError):
                pass

        return InferenceResult(
            content=content,
            parsed=parsed,
            endpoint_id=self._target.endpoint_id,
            deployment_id=self._target.deployment_id,
            model=metadata.get("model", self._target.model) or self._target.model,
            protocol=self._target.protocol,
            structured_mode=structured_mode,
            latency_ms=latency_ms,
            usage=TokenUsage(
                input_tokens=metadata.get("prompt_eval_count"),
                output_tokens=metadata.get("eval_count"),
            ),
            request_id=request.trace_id or None,
            repaired=False,
            retries=0,
        )

    def _build_payload(self, request: StructuredGenerationRequest) -> dict:
        """Build the Ollama /api/chat request body."""
        messages = [
            {"role": msg.role, "content": msg.content}
            for msg in request.messages
        ]

        options: dict = {}

        # Honor max output tokens via num_predict
        if request.max_output_tokens:
            options["num_predict"] = request.max_output_tokens

        # Honor context window from target configuration
        if self._target.context_window and self._target.context_window > 0:
            options["num_ctx"] = self._target.context_window

        # Temperature
        options["temperature"] = request.temperature

        # Seed if provided
        if request.seed is not None:
            options["seed"] = request.seed

        payload: dict = {
            "model": self._target.model,
            "messages": messages,
            "stream": True,
            "options": options,
        }

        # Native schema formatting when supported
        if request.json_schema and self._target.capabilities.json_schema:
            payload["format"] = request.json_schema

        # Merge any extra body params from target
        if self._target.extra_body:
            for key, value in self._target.extra_body.items():
                if key not in payload:
                    payload[key] = value

        return payload

    def _determine_structured_mode(
        self, request: StructuredGenerationRequest
    ) -> str:
        """Determine which structured output mode was used."""
        if not request.json_schema:
            return "none"
        if self._target.capabilities.json_schema:
            return "json_schema"
        # Schema was requested but not natively supported — prompt only
        return "prompt_only"

    async def _stream_response(
        self,
        url: str,
        payload: dict,
        timeout_seconds: float,
    ) -> tuple[str, dict]:
        """Stream Ollama response, applying stall detection.

        Returns (content, metadata) where metadata contains token counts
        and timing from the final Ollama response chunk.
        """
        content_parts: list[str] = []
        metadata: dict = {}
        last_content_length = 0
        unchanged_count = 0
        last_check_time = time.monotonic()

        try:
            async with self._http.stream(
                "POST",
                url,
                json=payload,
                timeout=httpx.Timeout(timeout_seconds, read=timeout_seconds),
            ) as response:
                if response.status_code != 200:
                    # Read body for error detail
                    body = b""
                    async for chunk in response.aiter_bytes():
                        body += chunk
                    detail = body.decode("utf-8", errors="replace")[:500]
                    category = _map_http_status_to_category(response.status_code)
                    raise InferenceError(
                        category,
                        f"Ollama returned HTTP {response.status_code}",
                        provider_detail=detail,
                        status_code=response.status_code,
                    )

                async for line in response.aiter_lines():
                    if not line.strip():
                        continue

                    try:
                        chunk_data = json.loads(line)
                    except json.JSONDecodeError:
                        continue

                    # Extract content from message
                    message = chunk_data.get("message", {})
                    chunk_content = message.get("content", "")
                    if chunk_content:
                        content_parts.append(chunk_content)

                    # Check for completion
                    if chunk_data.get("done", False):
                        # Capture metadata from final chunk
                        metadata["model"] = chunk_data.get("model", "")
                        metadata["prompt_eval_count"] = chunk_data.get("prompt_eval_count")
                        metadata["eval_count"] = chunk_data.get("eval_count")
                        total_ns = chunk_data.get("total_duration")
                        if total_ns:
                            metadata["total_duration_ms"] = total_ns // 1_000_000
                        break

                    # Stall detection
                    if self._stall_policy.enabled:
                        now = time.monotonic()
                        if now - last_check_time >= self._stall_policy.check_interval_seconds:
                            current_content = "".join(content_parts)
                            current_length = len(current_content)

                            if current_length == last_content_length:
                                unchanged_count += 1
                            else:
                                # Check for loop pattern
                                if _detect_loop(
                                    current_content,
                                    self._stall_policy.loop_window,
                                    self._stall_policy.loop_threshold,
                                ):
                                    unchanged_count += 1
                                else:
                                    unchanged_count = 0

                            last_content_length = current_length
                            last_check_time = now

                            if unchanged_count >= self._stall_policy.max_unchanged_intervals:
                                logger.warning(
                                    "Stall detected after %d unchanged intervals, aborting",
                                    unchanged_count,
                                )
                                raise InferenceError(
                                    InferenceErrorCategory.STALL_DETECTED,
                                    f"Generation stalled after {unchanged_count} check intervals",
                                )

        except InferenceError:
            raise
        except httpx.TimeoutException as exc:
            raise InferenceError(
                InferenceErrorCategory.TIMEOUT,
                "Stream read timed out",
                provider_detail=str(exc),
            ) from exc
        except httpx.ConnectError as exc:
            raise InferenceError(
                InferenceErrorCategory.CONNECTION_REFUSED,
                "Connection refused during streaming",
                provider_detail=str(exc),
            ) from exc
        except httpx.HTTPError as exc:
            raise InferenceError(
                InferenceErrorCategory.CONNECTION_ERROR,
                "HTTP error during streaming",
                provider_detail=str(exc),
            ) from exc

        final_content = "".join(content_parts)

        if not final_content:
            raise InferenceError(
                InferenceErrorCategory.EMPTY_RESPONSE,
                "Ollama returned empty content",
            )

        return final_content, metadata