from __future__ import annotations

from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any, Dict, List

import pytest

from entropic_core.core.llm_middleware import LLMMiddleware, LLMProvider, MiddlewareConfig


@dataclass
class DummyMessage:
    content: str


def _mk_openai_style_response(text: str) -> Any:
    # Mimic openai response shape: resp.choices[0].message.content
    return SimpleNamespace(
        choices=[SimpleNamespace(message=SimpleNamespace(content=text))],
        usage=SimpleNamespace(prompt_tokens=1, completion_tokens=2, total_tokens=3),
    )


def test_pre_call_injections_add_system_messages() -> None:
    cfg = MiddlewareConfig(
        enable_prompt_injection=True,
        enable_hallucination_detection=True,
        enable_retry_logic=False,
        enable_context_pruning=False,
        high_entropy_threshold=0.1,
        hallucination_threshold=0.1,
    )
    mw = LLMMiddleware(config=cfg)
    mw.update_entropy(0.9)

    seen: Dict[str, Any] = {}

    def create(*args, **kwargs):
        seen["messages"] = kwargs.get("messages", [])
        return _mk_openai_style_response("ok response with enough length to pass quality checks.")

    wrapped = mw.wrap(create, provider=LLMProvider.OPENAI)

    wrapped(model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], temperature=0.7)
    msgs: List[Dict[str, str]] = seen["messages"]
    system_text = "\n".join(m.get("content", "") for m in msgs if m.get("role") == "system")
    assert "[SYSTEM STABILITY PROTOCOL ACTIVE]" in system_text
    assert "[FACTUAL GROUNDING REQUIRED]" in system_text
    assert mw.call_history
    assert mw.call_history[-1].interventions_applied


def test_retry_logic_reduces_temperature_and_retries() -> None:
    cfg = MiddlewareConfig(
        enable_retry_logic=True,
        max_retries=2,
        retry_temperature_decay=0.5,
        enable_prompt_injection=False,
        enable_hallucination_detection=False,
        enable_context_pruning=False,
    )
    mw = LLMMiddleware(config=cfg)

    calls: List[float] = []

    def create(*args, **kwargs):
        calls.append(float(kwargs.get("temperature", 0.7)))
        # First response is "bad" (too short, repetitive, hedged); second response is longer (good).
        if len(calls) == 1:
            return _mk_openai_style_response("maybe maybe maybe maybe maybe maybe")
        return _mk_openai_style_response(
            "This is a sufficiently long response that should score as good quality."
        )

    wrapped = mw.wrap(create, provider=LLMProvider.OPENAI)
    wrapped(model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], temperature=0.8)

    assert len(calls) >= 2
    # Temperature decays on retry
    assert calls[1] == pytest.approx(calls[0] * 0.5)


def test_tool_messages_do_not_break_openai_adapter() -> None:
    cfg = MiddlewareConfig(
        enable_retry_logic=False,
        enable_prompt_injection=False,
        enable_hallucination_detection=False,
        enable_context_pruning=False,
    )
    mw = LLMMiddleware(config=cfg)

    def create(*args, **kwargs):
        # Just echo back to ensure the adapter can handle tool-role messages.
        return _mk_openai_style_response("ok " + str(len(kwargs.get("messages", []))))

    wrapped = mw.wrap(create, provider=LLMProvider.OPENAI)
    resp = wrapped(
        model="gpt-4o-mini",
        messages=[
            {"role": "user", "content": "call a tool"},
            {"role": "tool", "content": "tool output: 2"},
        ],
        temperature=0.7,
    )
    assert "ok" in resp.choices[0].message.content


def test_transport_retries_on_rate_limit_like_exception() -> None:
    class RateLimitLike(Exception):
        def __init__(self, status_code: int):
            super().__init__("rate limit")
            self.status_code = status_code

    cfg = MiddlewareConfig(
        enable_retry_logic=False,
        enable_prompt_injection=False,
        enable_hallucination_detection=False,
        enable_context_pruning=False,
        enable_transport_retries=True,
        transport_max_retries=2,
        transport_backoff_base_s=0.0,  # no sleep in tests
        transport_jitter_s=0.0,
    )
    mw = LLMMiddleware(config=cfg)

    calls = {"n": 0}

    def create(*args, **kwargs):
        calls["n"] += 1
        if calls["n"] == 1:
            raise RateLimitLike(429)
        return _mk_openai_style_response(
            "This is a sufficiently long response that should pass quality checks."
        )

    wrapped = mw.wrap(create, provider=LLMProvider.OPENAI)
    wrapped(model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], temperature=0.7)
    assert calls["n"] == 2
    assert mw.call_history


def test_streaming_wrap_preserves_iterator_and_records_history() -> None:
    # Fake OpenAI streaming chunks.
    def mk_chunk(delta_text: str):
        return SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content=delta_text))])

    def create(*args, **kwargs):
        assert kwargs.get("stream") is True

        def it():
            yield mk_chunk("Hello ")
            yield mk_chunk("world")

        return it()

    cfg = MiddlewareConfig(
        enable_retry_logic=False,
        enable_prompt_injection=False,
        enable_hallucination_detection=False,
        enable_context_pruning=False,
        enable_transport_retries=False,
    )
    mw = LLMMiddleware(config=cfg)

    wrapped = mw.wrap(create, provider=LLMProvider.OPENAI)
    stream = wrapped(model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=True)

    # Consume stream.
    got = list(stream)
    assert len(got) == 2
    # History recorded after consumption.
    assert mw.call_history
    assert mw.call_history[-1].is_stream is True
    assert "Hello world" in mw.call_history[-1].response_text


def test_idempotent_injections_do_not_duplicate_system_messages() -> None:
    cfg = MiddlewareConfig(
        enable_prompt_injection=True,
        enable_hallucination_detection=True,
        enable_retry_logic=False,
        enable_context_pruning=False,
        high_entropy_threshold=0.1,
        hallucination_threshold=0.1,
    )
    mw = LLMMiddleware(config=cfg)
    mw.update_entropy(0.9)

    adapter = mw._adapters[LLMProvider.OPENAI]
    req = {
        "model": "gpt-4o-mini",
        "messages": [{"role": "user", "content": "hi"}],
        "temperature": 0.7,
    }
    r1 = mw._apply_pre_call_interventions(req, adapter)
    r2 = mw._apply_pre_call_interventions(r1, adapter)

    msgs = adapter.extract_messages(r2)
    system_msgs = [m for m in msgs if m.get("role") == "system"]
    stability = [
        m for m in system_msgs if "[SYSTEM STABILITY PROTOCOL ACTIVE]" in str(m.get("content", ""))
    ]
    grounding = [
        m for m in system_msgs if "[FACTUAL GROUNDING REQUIRED]" in str(m.get("content", ""))
    ]
    assert len(stability) == 1
    assert len(grounding) == 1


def test_context_pruning_safe_never_removes_tool_messages() -> None:
    cfg = MiddlewareConfig(
        enable_retry_logic=False,
        enable_prompt_injection=False,
        enable_hallucination_detection=False,
        enable_context_pruning=True,
        max_context_tokens=10,  # force prune
    )
    mw = LLMMiddleware(config=cfg)
    adapter = mw._adapters[LLMProvider.OPENAI]

    # Put a tool message early so naive pruning might drop it.
    req = {
        "model": "gpt-4o-mini",
        "messages": [
            {"role": "system", "content": "sys"},
            {"role": "tool", "content": "tool output: IMPORTANT E1"},
            {"role": "user", "content": "x" * 200},
            {"role": "assistant", "content": "y" * 200},
            {"role": "user", "content": "z" * 200},
        ],
    }
    out = mw._apply_pre_call_interventions(req, adapter)
    msgs = adapter.extract_messages(out)
    assert any(m.get("role") == "tool" for m in msgs)
