from __future__ import annotations

from types import SimpleNamespace

import entropic_core
from entropic_core.core.llm_middleware import MiddlewareConfig


def _mk_openai_style_response(text: str):
    # Mimic openai response shape: resp.choices[0].message.content
    return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content=text))])


def fake_llm_create(*args, **kwargs):
    # Echo back to show what the middleware injected.
    messages = kwargs.get("messages", [])
    sys_count = sum(1 for m in messages if m.get("role") == "system")
    tool_count = sum(1 for m in messages if m.get("role") == "tool")
    return _mk_openai_style_response(f"ok (system={sys_count}, tool={tool_count}) " + ("x" * 80))


def main() -> int:
    cfg = MiddlewareConfig(
        enable_retry_logic=False,
        enable_context_pruning=True,
        max_context_tokens=10,  # force pruning
        enable_prompt_injection=True,
        enable_hallucination_detection=True,
        high_entropy_threshold=0.1,
        hallucination_threshold=0.1,
    )
    mw = entropic_core.create_middleware(cfg)
    mw.update_entropy(0.9)

    wrapped = mw.wrap(fake_llm_create)

    req = {
        "model": "dummy",
        "messages": [
            {"role": "system", "content": "sys"},
            {"role": "tool", "content": "tool output: IMPORTANT E1"},
            {"role": "user", "content": "hello " + ("x" * 400)},
            {"role": "assistant", "content": "previous " + ("y" * 400)},
            {"role": "user", "content": "question " + ("z" * 400)},
        ],
        "temperature": 0.7,
    }

    # Call through a provider adapter by using the OpenAI-ish signature.
    resp = wrapped(model=req["model"], messages=req["messages"], temperature=req["temperature"])
    print(resp.choices[0].message.content)
    print(mw.get_stats())
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
