"""Tests for OllamaNativeClient — shared inference gateway Ollama implementation.

Covers:
- Successful generation with native schema format
- Prompt-only mode reporting when schema not supported
- Max tokens and context window configuration
- Stall detection triggers abort
- Error mapping (connection refused, timeout, model not found)
- Result captures correct metadata (model, duration, token counts)
- Stall detection is Ollama-specific (doesn't affect InferenceResult interface)

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

import json
from uuid import uuid4

import httpx
import pytest

from services.shared.inference.clients.ollama_native import (
    OllamaNativeClient,
    StallPolicy,
    _detect_loop,
    _map_http_status_to_category,
)
from services.shared.inference.errors import InferenceError, InferenceErrorCategory
from services.shared.inference.models import (
    ChatMessage,
    InferenceResult,
    InferenceTarget,
    ProviderCapabilities,
    StructuredGenerationRequest,
)


def _make_target(
    *,
    json_schema_capable: bool = True,
    context_window: int | None = None,
    max_output_tokens: int | None = None,
) -> InferenceTarget:
    """Create a test InferenceTarget for Ollama."""
    return InferenceTarget(
        endpoint_id=uuid4(),
        deployment_id=uuid4(),
        protocol="ollama_native",
        base_url="http://test-ollama:11434",
        model="test-model:7b",
        capabilities=ProviderCapabilities(
            chat_completions=True,
            json_schema=json_schema_capable,
        ),
        context_window=context_window,
        max_output_tokens=max_output_tokens,
    )


def _make_request(
    *,
    schema: dict | None = None,
    max_output_tokens: int = 2048,
    temperature: float = 0.0,
    seed: int | None = 0,
) -> StructuredGenerationRequest:
    """Create a test StructuredGenerationRequest."""
    return StructuredGenerationRequest(
        messages=[
            ChatMessage(role="system", content="You are a financial analyst."),
            ChatMessage(role="user", content="Analyze AAPL earnings."),
        ],
        json_schema=schema,
        max_output_tokens=max_output_tokens,
        temperature=temperature,
        seed=seed,
        timeout_seconds=30.0,
        trace_id="test-trace-001",
    )


def _streaming_response(
    content: str,
    *,
    model: str = "test-model:7b",
    prompt_eval_count: int = 150,
    eval_count: int = 200,
    total_duration_ns: int = 5_000_000_000,
) -> list[str]:
    """Build Ollama streaming response lines (newline-delimited JSON)."""
    lines = []
    # Stream content in chunks
    chunk_size = max(1, len(content) // 3)
    for i in range(0, len(content), chunk_size):
        chunk = content[i:i + chunk_size]
        lines.append(json.dumps({
            "model": model,
            "message": {"role": "assistant", "content": chunk},
            "done": False,
        }))

    # Final done message with metadata
    lines.append(json.dumps({
        "model": model,
        "message": {"role": "assistant", "content": ""},
        "done": True,
        "prompt_eval_count": prompt_eval_count,
        "eval_count": eval_count,
        "total_duration": total_duration_ns,
    }))
    return lines


def _make_streaming_transport(
    lines: list[str],
    *,
    status_code: int = 200,
) -> httpx.MockTransport:
    """Build a mock transport that returns streaming lines."""
    body = "\n".join(lines) + "\n"

    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            status_code,
            content=body.encode(),
            headers={"content-type": "application/x-ndjson"},
        )

    return httpx.MockTransport(handler)


# --- Test: Successful generation with native schema format ---


@pytest.mark.asyncio
async def test_generate_success_with_schema_format():
    """Successful generation uses native format field when json_schema capable."""
    result_json = json.dumps({"signal": "bullish", "confidence": 0.85})
    lines = _streaming_response(result_json, eval_count=50, prompt_eval_count=100)
    transport = _make_streaming_transport(lines)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target(json_schema_capable=True)
    client = OllamaNativeClient(target, http_client=http)

    schema = {"type": "object", "properties": {"signal": {"type": "string"}}}
    request = _make_request(schema=schema)

    result = await client.generate(request)

    assert isinstance(result, InferenceResult)
    assert result.structured_mode == "json_schema"
    assert result.parsed == {"signal": "bullish", "confidence": 0.85}
    assert result.content == result_json
    assert result.model == "test-model:7b"
    assert result.usage.output_tokens == 50
    assert result.usage.input_tokens == 100
    assert result.request_id == "test-trace-001"
    assert result.latency_ms >= 0
    assert result.endpoint_id == target.endpoint_id
    assert result.deployment_id == target.deployment_id
    assert result.protocol == "ollama_native"

    await client.close()


# --- Test: Prompt-only mode reporting ---


@pytest.mark.asyncio
async def test_generate_prompt_only_mode():
    """Reports prompt_only when schema requested but not natively supported."""
    result_json = json.dumps({"signal": "bearish"})
    lines = _streaming_response(result_json)
    transport = _make_streaming_transport(lines)
    http = httpx.AsyncClient(transport=transport)

    # No json_schema capability
    target = _make_target(json_schema_capable=False)
    client = OllamaNativeClient(target, http_client=http)

    schema = {"type": "object", "properties": {"signal": {"type": "string"}}}
    request = _make_request(schema=schema)

    result = await client.generate(request)

    assert result.structured_mode == "prompt_only"
    assert result.parsed == {"signal": "bearish"}

    await client.close()


@pytest.mark.asyncio
async def test_generate_no_schema_mode_none():
    """Reports 'none' when no schema is requested."""
    lines = _streaming_response("Free text response about markets.")
    transport = _make_streaming_transport(lines)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target(json_schema_capable=True)
    client = OllamaNativeClient(target, http_client=http)

    request = _make_request(schema=None)

    result = await client.generate(request)

    assert result.structured_mode == "none"
    assert result.parsed is None
    assert "Free text response" in result.content

    await client.close()


# --- Test: Max tokens and context configuration ---


@pytest.mark.asyncio
async def test_max_tokens_and_context_in_payload():
    """num_predict and num_ctx are set in Ollama options from request/target."""
    captured_payload: dict = {}

    def handler(request: httpx.Request) -> httpx.Response:
        captured_payload.update(json.loads(request.content))
        body = "\n".join(_streaming_response("ok"))
        return httpx.Response(200, content=body.encode())

    transport = httpx.MockTransport(handler)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target(context_window=32768)
    client = OllamaNativeClient(target, http_client=http)

    request = _make_request(max_output_tokens=4096, temperature=0.1, seed=42)

    await client.generate(request)

    assert captured_payload["options"]["num_predict"] == 4096
    assert captured_payload["options"]["num_ctx"] == 32768
    assert captured_payload["options"]["temperature"] == 0.1
    assert captured_payload["options"]["seed"] == 42
    assert captured_payload["model"] == "test-model:7b"
    assert captured_payload["stream"] is True

    await client.close()


@pytest.mark.asyncio
async def test_schema_passed_as_format_field():
    """When json_schema is capable, the schema is passed via format field."""
    captured_payload: dict = {}

    def handler(request: httpx.Request) -> httpx.Response:
        captured_payload.update(json.loads(request.content))
        body = "\n".join(_streaming_response('{"x": 1}'))
        return httpx.Response(200, content=body.encode())

    transport = httpx.MockTransport(handler)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target(json_schema_capable=True)
    client = OllamaNativeClient(target, http_client=http)

    schema = {"type": "object", "properties": {"x": {"type": "integer"}}}
    request = _make_request(schema=schema)

    await client.generate(request)

    assert captured_payload["format"] == schema

    await client.close()


@pytest.mark.asyncio
async def test_no_format_field_when_not_capable():
    """When json_schema is not capable, format field is omitted."""
    captured_payload: dict = {}

    def handler(request: httpx.Request) -> httpx.Response:
        captured_payload.update(json.loads(request.content))
        body = "\n".join(_streaming_response('{"x": 1}'))
        return httpx.Response(200, content=body.encode())

    transport = httpx.MockTransport(handler)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target(json_schema_capable=False)
    client = OllamaNativeClient(target, http_client=http)

    schema = {"type": "object", "properties": {"x": {"type": "integer"}}}
    request = _make_request(schema=schema)

    await client.generate(request)

    assert "format" not in captured_payload

    await client.close()


# --- Test: Stall detection triggers abort ---


@pytest.mark.asyncio
async def test_stall_detection_triggers_abort():
    """Stall detection raises InferenceError with STALL_DETECTED category."""
    # Generate highly repetitive content that triggers loop detection
    repeated = "buy buy buy " * 200  # Very repetitive

    # Build streaming lines that produce repetitive content slowly
    lines = []
    chunk_size = 50
    for i in range(0, len(repeated), chunk_size):
        chunk = repeated[i:i + chunk_size]
        lines.append(json.dumps({
            "model": "test-model:7b",
            "message": {"role": "assistant", "content": chunk},
            "done": False,
        }))
    # Add done at the end (but stall should abort before reaching it)
    lines.append(json.dumps({
        "model": "test-model:7b",
        "message": {"role": "assistant", "content": ""},
        "done": True,
        "eval_count": 500,
    }))

    transport = _make_streaming_transport(lines)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target()
    # Aggressive stall policy for testing
    stall_policy = StallPolicy(
        enabled=True,
        check_interval_seconds=0.0,  # Check every iteration
        max_unchanged_intervals=2,
        loop_window=48,
        loop_threshold=0.5,
    )
    client = OllamaNativeClient(target, http_client=http, stall_policy=stall_policy)

    request = _make_request()

    with pytest.raises(InferenceError) as exc_info:
        await client.generate(request)

    assert exc_info.value.category == InferenceErrorCategory.STALL_DETECTED
    assert exc_info.value.retryable is True

    await client.close()


@pytest.mark.asyncio
async def test_stall_detection_disabled():
    """When stall detection is disabled, repetitive content completes normally."""
    repeated = "x" * 500
    lines = _streaming_response(repeated)
    transport = _make_streaming_transport(lines)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target()
    stall_policy = StallPolicy(enabled=False)
    client = OllamaNativeClient(target, http_client=http, stall_policy=stall_policy)

    request = _make_request()

    result = await client.generate(request)
    assert result.content == repeated

    await client.close()


# --- Test: Error mapping ---


@pytest.mark.asyncio
async def test_error_connection_refused():
    """Connection refused maps to CONNECTION_REFUSED category."""

    def handler(request: httpx.Request) -> httpx.Response:
        raise httpx.ConnectError("Connection refused")

    transport = httpx.MockTransport(handler)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target()
    client = OllamaNativeClient(target, http_client=http)
    request = _make_request()

    with pytest.raises(InferenceError) as exc_info:
        await client.generate(request)

    assert exc_info.value.category == InferenceErrorCategory.CONNECTION_REFUSED
    assert exc_info.value.retryable is True

    await client.close()


@pytest.mark.asyncio
async def test_error_timeout():
    """Timeout maps to TIMEOUT category."""

    def handler(request: httpx.Request) -> httpx.Response:
        raise httpx.ReadTimeout("timed out")

    transport = httpx.MockTransport(handler)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target()
    client = OllamaNativeClient(target, http_client=http)
    request = _make_request()

    with pytest.raises(InferenceError) as exc_info:
        await client.generate(request)

    assert exc_info.value.category == InferenceErrorCategory.TIMEOUT
    assert exc_info.value.retryable is True

    await client.close()


@pytest.mark.asyncio
async def test_error_model_not_found():
    """HTTP 404 maps to MODEL_NOT_FOUND category."""
    error_body = json.dumps({"error": "model 'nonexistent' not found"})

    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(404, content=error_body.encode())

    transport = httpx.MockTransport(handler)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target()
    client = OllamaNativeClient(target, http_client=http)
    request = _make_request()

    with pytest.raises(InferenceError) as exc_info:
        await client.generate(request)

    assert exc_info.value.category == InferenceErrorCategory.MODEL_NOT_FOUND
    assert exc_info.value.status_code == 404
    assert exc_info.value.retryable is False

    await client.close()


@pytest.mark.asyncio
async def test_error_server_error():
    """HTTP 500 maps to SERVER_ERROR category."""
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(500, content=b"Internal Server Error")

    transport = httpx.MockTransport(handler)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target()
    client = OllamaNativeClient(target, http_client=http)
    request = _make_request()

    with pytest.raises(InferenceError) as exc_info:
        await client.generate(request)

    assert exc_info.value.category == InferenceErrorCategory.SERVER_ERROR
    assert exc_info.value.retryable is True

    await client.close()


@pytest.mark.asyncio
async def test_error_empty_response():
    """Empty content from Ollama raises EMPTY_RESPONSE."""
    # Only a done message with no content
    lines = [json.dumps({
        "model": "test-model:7b",
        "message": {"role": "assistant", "content": ""},
        "done": True,
        "eval_count": 0,
    })]
    transport = _make_streaming_transport(lines)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target()
    client = OllamaNativeClient(target, http_client=http)
    request = _make_request()

    with pytest.raises(InferenceError) as exc_info:
        await client.generate(request)

    assert exc_info.value.category == InferenceErrorCategory.EMPTY_RESPONSE

    await client.close()


# --- Test: Result captures correct metadata ---


@pytest.mark.asyncio
async def test_result_metadata_captured():
    """InferenceResult captures model name, duration, and token counts from Ollama metadata."""
    content = json.dumps({"result": "test"})
    lines = _streaming_response(
        content,
        model="qwen3.5:9b",
        prompt_eval_count=250,
        eval_count=180,
        total_duration_ns=8_500_000_000,
    )
    transport = _make_streaming_transport(lines)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target()
    client = OllamaNativeClient(target, http_client=http)
    request = _make_request()

    result = await client.generate(request)

    assert result.model == "qwen3.5:9b"
    assert result.usage.input_tokens == 250
    assert result.usage.output_tokens == 180
    assert result.latency_ms >= 0
    assert result.endpoint_id == target.endpoint_id
    assert result.deployment_id == target.deployment_id
    assert result.protocol == "ollama_native"

    await client.close()


# --- Test: Stall detection is Ollama-specific, doesn't affect interface ---


@pytest.mark.asyncio
async def test_stall_detection_ollama_specific_not_in_result():
    """Stall detection is Ollama-specific policy — successful results have no stall fields."""
    content = "Normal varied content with different words and ideas flowing naturally."
    lines = _streaming_response(content)
    transport = _make_streaming_transport(lines)
    http = httpx.AsyncClient(transport=transport)

    target = _make_target()
    stall_policy = StallPolicy(enabled=True)
    client = OllamaNativeClient(target, http_client=http, stall_policy=stall_policy)
    request = _make_request()

    result = await client.generate(request)

    # InferenceResult has no stall-related fields — it's protocol agnostic
    assert isinstance(result, InferenceResult)
    result_fields = set(InferenceResult.model_fields.keys())
    assert "stall_detected" not in result_fields
    assert "stall_policy" not in result_fields
    # The result is normal
    assert result.content == content

    await client.close()


# --- Test: Protocol validation ---


def test_wrong_protocol_raises():
    """OllamaNativeClient rejects non-ollama_native targets."""
    target = InferenceTarget(
        endpoint_id=uuid4(),
        deployment_id=uuid4(),
        protocol="openai_chat",
        base_url="http://test:8000",
        model="gpt-4",
        capabilities=ProviderCapabilities(),
    )

    with pytest.raises(ValueError, match="ollama_native"):
        OllamaNativeClient(target)


# --- Test: Helper functions ---


def test_detect_loop_with_repetition():
    """Loop detection catches repeated tail content."""
    content = "hello world " * 20
    assert _detect_loop(content, window=48, threshold=0.5) is True


def test_detect_loop_varied_content():
    """Loop detection does not trigger on varied content with default threshold."""
    content = (
        "Apple reported strong Q4 earnings with revenue up 12 percent. "
        "The company saw growth across all segments including services and Mac. "
        "CEO Tim Cook highlighted the success of Apple Intelligence features. "
        "Analysts raised their price targets following the announcement. "
        "Markets rallied today on strong earnings from tech sector leaders."
    )
    # Default threshold is 0.5 — normal English text has ~0.30 unique ratio
    # which is above 0.15 but below 0.5 — the key is tail not in body
    assert _detect_loop(content, window=64, threshold=0.15) is False


def test_detect_loop_short_content():
    """Loop detection returns False for content shorter than 2x window."""
    content = "short"
    assert _detect_loop(content, window=64, threshold=0.5) is False


def test_map_http_status_categories():
    """HTTP status codes map to correct error categories."""
    assert _map_http_status_to_category(401) == InferenceErrorCategory.AUTH_FAILED
    assert _map_http_status_to_category(403) == InferenceErrorCategory.FORBIDDEN
    assert _map_http_status_to_category(404) == InferenceErrorCategory.MODEL_NOT_FOUND
    assert _map_http_status_to_category(429) == InferenceErrorCategory.RATE_LIMITED
    assert _map_http_status_to_category(400) == InferenceErrorCategory.BAD_REQUEST
    assert _map_http_status_to_category(500) == InferenceErrorCategory.SERVER_ERROR
    assert _map_http_status_to_category(503) == InferenceErrorCategory.SERVER_ERROR
    assert _map_http_status_to_category(418) == InferenceErrorCategory.UNKNOWN