"""
Integration test for completions endpoint token threshold filtering.

This test verifies that the /api/v1/chat/completions endpoint properly
filters session history to stay within both message count and token limits.

Two scenarios:
1. High message count: Many small messages (exceeds message count limit)
2. High token count: Few large messages (exceeds token threshold)

The bug: completions.py loads raw_session_history WITHOUT applying
filter_within_token_threshold, causing context_length_exceeded errors
for sessions with large history.
"""

import os
from pathlib import Path

import pytest
import uuid
import yaml
from datetime import datetime, timezone, timedelta

from rem.services.session import SessionMessageStore
from rem.services.session.compression import filter_within_token_threshold
from rem.settings import settings


# Load fixture data
FIXTURE_PATH = Path(__file__).parent.parent / "fixtures" / "large_session_fixture.yaml"


def load_fixture():
    """Load the large session fixture."""
    with open(FIXTURE_PATH) as f:
        return yaml.safe_load(f)


@pytest.fixture
def session_id():
    return str(uuid.uuid4())


@pytest.fixture
def user_id():
    return str(uuid.uuid4())


async def create_session_from_scenario(
    session_id: str,
    user_id: str,
    scenario: dict,
) -> list[dict]:
    """
    Create a session with messages from a fixture scenario.

    Args:
        session_id: Session ID
        user_id: User ID
        scenario: Fixture scenario dict with templates and generation config
    """
    store = SessionMessageStore(user_id=user_id)

    user_template = scenario["user_message_template"]
    assistant_template = scenario["assistant_message_template"]
    pairs = scenario["generation"]["pairs"]
    time_spread_days = scenario["generation"]["time_spread_days"]

    messages = []
    base_time = datetime.now(timezone.utc) - timedelta(days=time_spread_days)

    for i in range(pairs):
        # User message from template
        user_content = user_template.format(index=i)
        messages.append({
            "role": "user",
            "content": user_content,
            "timestamp": (base_time + timedelta(hours=i*2)).isoformat(),
        })

        # Assistant response from template
        assistant_content = assistant_template.format(index=i)
        messages.append({
            "role": "assistant",
            "content": assistant_content,
            "timestamp": (base_time + timedelta(hours=i*2 + 1)).isoformat(),
        })

    # Store all messages
    await store.store_session_messages(
        session_id=session_id,
        messages=messages,
        user_id=user_id,
        compress=True,
    )

    return messages


class TestCompletionsTokenThreshold:
    """
    Integration tests verifying token threshold is applied in completions endpoint.
    """

    @pytest.mark.asyncio
    async def test_high_message_count_exceeds_limit(self, session_id: str, user_id: str):
        """
        Test scenario 1: Many small messages exceeding message count limit.
        Verifies the raw session has more messages than allowed.
        """
        fixture = load_fixture()
        scenario = fixture["high_message_count"]

        messages = await create_session_from_scenario(session_id, user_id, scenario)

        store = SessionMessageStore(user_id=user_id)
        raw_history, _ = await store.load_session_messages(
            session_id=session_id,
            user_id=user_id,
            compress_on_load=False,
        )

        # Should have many messages
        expected_messages = scenario["generation"]["pairs"] * 2
        assert len(raw_history) == expected_messages, (
            f"Expected {expected_messages} messages, got {len(raw_history)}"
        )

        print(f"High message count scenario: {len(raw_history)} messages")

    @pytest.mark.asyncio
    async def test_high_token_count_setup(self, session_id: str, user_id: str):
        """
        Test scenario 2: Few large messages exceeding token threshold.
        Verifies the raw session exceeds token threshold despite fewer messages.
        """
        fixture = load_fixture()
        scenario = fixture["high_token_count"]

        messages = await create_session_from_scenario(session_id, user_id, scenario)

        store = SessionMessageStore(user_id=user_id)
        raw_history, _ = await store.load_session_messages(
            session_id=session_id,
            user_id=user_id,
            compress_on_load=False,
        )

        # Count tokens (rough estimate: 4 chars per token)
        total_chars = sum(len(m.get("content", "")) for m in raw_history)
        estimated_tokens = total_chars // 4

        # Should exceed token threshold
        assert estimated_tokens > settings.moment_builder.token_threshold, (
            f"Test data ({estimated_tokens} tokens) should exceed threshold "
            f"({settings.moment_builder.token_threshold})"
        )

        # But should have fewer messages than the high message count scenario
        expected_messages = scenario["generation"]["pairs"] * 2
        assert len(raw_history) == expected_messages
        assert expected_messages <= 250, "Token scenario should have relatively few messages"

        print(f"High token count scenario: {len(raw_history)} messages, ~{estimated_tokens} tokens")

    @pytest.mark.asyncio
    async def test_filter_reduces_high_token_session(self, session_id: str, user_id: str):
        """
        Verify filter_within_token_threshold works correctly on high token sessions.
        """
        fixture = load_fixture()
        scenario = fixture["high_token_count"]

        await create_session_from_scenario(session_id, user_id, scenario)

        store = SessionMessageStore(user_id=user_id)
        raw_history, _ = await store.load_session_messages(
            session_id=session_id,
            user_id=user_id,
            compress_on_load=False,
        )

        # Apply the filter
        filtered, total_tokens = await filter_within_token_threshold(
            messages=raw_history,
            session_id=session_id,
            user_id=user_id,
            token_threshold=settings.moment_builder.token_threshold,
            model="gpt-4.1",
        )

        # Should be reduced
        assert total_tokens <= settings.moment_builder.token_threshold, (
            f"Filtered tokens ({total_tokens}) should be <= threshold "
            f"({settings.moment_builder.token_threshold})"
        )

        print(f"Filtered: {len(filtered)} messages from {len(raw_history)}, {total_tokens} tokens")

    @pytest.mark.asyncio
    @pytest.mark.llm
    async def test_completions_endpoint_applies_token_filter(self, session_id: str, user_id: str):
        """
        CRITICAL TEST: Verify completions endpoint applies token threshold filter.

        This tests the bug where completions.py doesn't filter large sessions.
        Uses the 'high_token_count' scenario with few but large messages.

        Marked with @pytest.mark.llm - requires real API keys to run.
        Run with: pytest -m llm ...
        Skip with: pytest -m "not llm" ...
        """
        fixture = load_fixture()
        scenario = fixture["high_token_count"]

        await create_session_from_scenario(session_id, user_id, scenario)

        from httpx import AsyncClient, ASGITransport
        from rem.api.main import app

        async with AsyncClient(
            transport=ASGITransport(app=app),
            base_url="http://test"
        ) as client:
            response = await client.post(
                "/api/v1/chat/completions",
                json={
                    "model": "openai:gpt-4.1",
                    "messages": [{"role": "user", "content": "Summarize our conversation."}],
                    "stream": False,
                },
                headers={
                    "X-User-Id": user_id,
                    "X-Session-Id": session_id,
                    "X-Agent-Schema": "rem",
                    "Content-Type": "application/json",
                },
                timeout=120.0,
            )

        # Should NOT fail with context_length_exceeded
        # Before fix: fails with 500 (context_length_exceeded from LLM)
        # After fix: succeeds with 200
        assert response.status_code == 200, (
            f"Completions endpoint should succeed but got {response.status_code}: "
            f"{response.text}"
        )

        data = response.json()
        assert "choices" in data
        assert len(data["choices"]) > 0
        assert data["choices"][0]["message"]["content"]

        print(f"Completions succeeded: {data['choices'][0]['message']['content'][:100]}...")

    @pytest.mark.asyncio
    @pytest.mark.llm
    async def test_completions_streaming_applies_token_filter(self, session_id: str, user_id: str):
        """
        Same as above but for streaming mode.

        Marked with @pytest.mark.llm - requires real API keys to run.
        """
        fixture = load_fixture()
        scenario = fixture["high_token_count"]

        await create_session_from_scenario(session_id, user_id, scenario)

        from httpx import AsyncClient, ASGITransport
        from rem.api.main import app

        async with AsyncClient(
            transport=ASGITransport(app=app),
            base_url="http://test"
        ) as client:
            response = await client.post(
                "/api/v1/chat/completions",
                json={
                    "model": "openai:gpt-4.1",
                    "messages": [{"role": "user", "content": "Summarize our conversation."}],
                    "stream": True,
                },
                headers={
                    "X-User-Id": user_id,
                    "X-Session-Id": session_id,
                    "X-Agent-Schema": "rem",
                    "Content-Type": "application/json",
                },
                timeout=120.0,
            )

        # Should NOT fail with context_length_exceeded
        assert response.status_code == 200, (
            f"Streaming completions should succeed but got {response.status_code}: "
            f"{response.text}"
        )
