"""
Tests for Brynq Chat API — POST /api/chat, POST /api/chat/chain,
GET /api/models, GET /api/ollama/status, and WebSocket /ws/chat.
"""

import asyncio
import json
import time
from unittest.mock import MagicMock, patch, AsyncMock

import pytest


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def mock_llm_router():
    """Patch llm_router functions used by chat endpoints."""
    with patch("brynq.llm_router.chat", return_value="Mock LLM response") as mock_chat, \
         patch("brynq.llm_router._pick_first_available", return_value=("ollama", "llama3")), \
         patch("brynq.llm_router._default_model_for", return_value="llama3"), \
         patch("brynq.llm_router.discover_backends", return_value=[
             {"name": "ollama", "status": "running", "models": ["llama3", "mistral"], "url": "localhost:11434"},
             {"name": "lmstudio", "status": "unavailable", "models": [], "url": ""},
             {"name": "claude", "status": "not_configured", "models": ["claude-sonnet-4-20250514"], "url": "api.anthropic.com"},
             {"name": "gemini", "status": "configured", "models": ["gemini-2.5-pro"], "url": "generativelanguage.googleapis.com"},
         ]), \
         patch("brynq.llm_router._probe_ollama", return_value={
             "name": "ollama", "status": "running", "models": ["llama3", "mistral"], "url": "localhost:11434",
         }):
        yield mock_chat


@pytest.fixture
def test_client(mock_llm_router):
    """Create a FastAPI TestClient with mocked LLM router."""
    from fastapi.testclient import TestClient
    from brynq.api_server import app
    return TestClient(app)


# ---------------------------------------------------------------------------
# POST /api/chat
# ---------------------------------------------------------------------------

class TestChatEndpoint:
    def test_basic_chat(self, test_client):
        resp = test_client.post("/api/chat", json={
            "message": "What is Q4 revenue?",
            "model": "llama3",
            "backend": "ollama",
        })
        assert resp.status_code == 200
        data = resp.json()
        assert "reply" in data
        assert data["reply"] == "Mock LLM response"
        assert data["model"] == "llama3"
        assert data["backend"] == "ollama"
        assert "elapsed_ms" in data
        assert isinstance(data["elapsed_ms"], int)

    def test_chat_auto_select(self, test_client):
        resp = test_client.post("/api/chat", json={
            "message": "Hello",
        })
        assert resp.status_code == 200
        data = resp.json()
        assert data["reply"] == "Mock LLM response"
        # Should auto-select ollama/llama3
        assert data["backend"] == "ollama"
        assert data["model"] == "llama3"

    def test_chat_with_system_prompt(self, test_client, mock_llm_router):
        resp = test_client.post("/api/chat", json={
            "message": "What is 2+2?",
            "system_prompt": "You are a math tutor.",
            "model": "llama3",
            "backend": "ollama",
        })
        assert resp.status_code == 200
        # Verify system prompt was included in the messages
        call_args = mock_llm_router.call_args
        messages = call_args.kwargs.get("messages") or call_args[0][2]
        assert any(m["role"] == "system" for m in messages)

    def test_chat_with_conversation_history(self, test_client, mock_llm_router):
        resp = test_client.post("/api/chat", json={
            "message": "And what about Q3?",
            "model": "llama3",
            "backend": "ollama",
            "conversation": [
                {"role": "user", "content": "What is Q4 revenue?"},
                {"role": "assistant", "content": "Q4 revenue was $2.3M"},
            ],
        })
        assert resp.status_code == 200
        call_args = mock_llm_router.call_args
        messages = call_args.kwargs.get("messages") or call_args[0][2]
        # Should have the conversation history + new message
        assert len(messages) >= 3

    def test_chat_with_temperature(self, test_client, mock_llm_router):
        resp = test_client.post("/api/chat", json={
            "message": "Be creative",
            "model": "llama3",
            "backend": "ollama",
            "temperature": 0.9,
        })
        assert resp.status_code == 200
        call_kwargs = mock_llm_router.call_args.kwargs
        assert call_kwargs.get("temperature") == 0.9

    def test_chat_empty_message_rejected(self, test_client):
        resp = test_client.post("/api/chat", json={
            "message": "",
        })
        assert resp.status_code == 422  # Pydantic validation error

    def test_chat_sse_streaming(self, test_client):
        resp = test_client.post(
            "/api/chat",
            json={"message": "Stream this", "model": "llama3", "backend": "ollama"},
            headers={"Accept": "text/event-stream"},
        )
        assert resp.status_code == 200
        assert "text/event-stream" in resp.headers.get("content-type", "")

        # Parse SSE events
        events = []
        for line in resp.text.strip().split("\n"):
            if line.startswith("data: "):
                events.append(json.loads(line[6:]))

        assert len(events) == 2
        assert events[0]["type"] == "chunk"
        assert events[0]["content"] == "Mock LLM response"
        assert events[1]["type"] == "done"
        assert "elapsed_ms" in events[1]


# ---------------------------------------------------------------------------
# POST /api/chat/chain
# ---------------------------------------------------------------------------

class TestChatChainEndpoint:
    def test_basic_chain(self, test_client, mock_llm_router):
        resp = test_client.post("/api/chat/chain", json={
            "message": "Analyze revenue",
            "steps": [
                {"backend": "ollama", "model": "llama3", "prompt": "Analyze: {original_input}"},
                {"backend": "ollama", "model": "llama3", "prompt": "Summarize: {step_1_output}"},
            ],
        })
        assert resp.status_code == 200
        data = resp.json()
        assert data["success"] is True
        assert "final_output" in data
        assert "total_elapsed_ms" in data
        assert len(data["steps"]) == 2
        for step in data["steps"]:
            assert "step_name" in step
            assert "backend" in step
            assert "model" in step
            assert "output" in step
            assert "elapsed_ms" in step
            assert "success" in step

    def test_chain_empty_steps_rejected(self, test_client):
        resp = test_client.post("/api/chat/chain", json={
            "message": "Test",
            "steps": [],
        })
        assert resp.status_code == 422

    def test_chain_step_failure(self, test_client, mock_llm_router):
        mock_llm_router.side_effect = RuntimeError("Model unavailable")
        resp = test_client.post("/api/chat/chain", json={
            "message": "Test",
            "steps": [
                {"backend": "ollama", "model": "llama3", "prompt": "{original_input}"},
            ],
        })
        assert resp.status_code == 200
        data = resp.json()
        assert data["success"] is False
        assert data["steps"][0]["success"] is False
        assert data["steps"][0]["error"] is not None

    def test_chain_missing_message(self, test_client):
        resp = test_client.post("/api/chat/chain", json={
            "steps": [{"prompt": "Hello"}],
        })
        assert resp.status_code == 422


# ---------------------------------------------------------------------------
# GET /api/models
# ---------------------------------------------------------------------------

class TestModelsEndpoint:
    def test_list_models(self, test_client):
        resp = test_client.get("/api/models")
        assert resp.status_code == 200
        data = resp.json()
        assert "models" in data
        models = data["models"]
        assert len(models) > 0

        # Check that local ollama models are listed
        ollama_models = [m for m in models if m["backend"] == "ollama"]
        assert len(ollama_models) >= 1
        assert ollama_models[0]["type"] == "local"
        assert ollama_models[0]["status"] == "ready"

    def test_models_include_cloud(self, test_client):
        resp = test_client.get("/api/models")
        data = resp.json()
        models = data["models"]

        cloud_models = [m for m in models if m["type"] == "cloud"]
        assert len(cloud_models) > 0

        # Claude should be no_key (not_configured in mock)
        claude_models = [m for m in models if m["backend"] == "claude"]
        assert len(claude_models) > 0
        assert claude_models[0]["status"] == "no_key"

        # Gemini should be key_configured (configured in mock)
        gemini_models = [m for m in models if m["backend"] == "gemini"]
        assert len(gemini_models) > 0
        assert gemini_models[0]["status"] == "key_configured"

    def test_models_schema(self, test_client):
        resp = test_client.get("/api/models")
        data = resp.json()
        for model in data["models"]:
            assert "name" in model
            assert "backend" in model
            assert "type" in model
            assert model["type"] in ("local", "cloud")
            assert "status" in model


# ---------------------------------------------------------------------------
# GET /api/ollama/status
# ---------------------------------------------------------------------------

class TestOllamaStatusEndpoint:
    def test_ollama_running(self, test_client):
        resp = test_client.get("/api/ollama/status")
        assert resp.status_code == 200
        data = resp.json()
        assert data["running"] is True
        assert data["status"] == "running"
        assert "models" in data
        assert "llama3" in data["models"]

    def test_ollama_not_running(self, test_client):
        with patch("brynq.llm_router._probe_ollama", return_value={
            "name": "ollama", "status": "unavailable", "models": [], "url": "",
        }):
            resp = test_client.get("/api/ollama/status")
            assert resp.status_code == 200
            data = resp.json()
            assert data["running"] is False
            assert data["status"] == "unavailable"


# ---------------------------------------------------------------------------
# WebSocket /ws/chat
# ---------------------------------------------------------------------------

class TestChatWebSocket:
    def test_ws_basic_message(self, test_client):
        with test_client.websocket_connect("/ws/chat") as ws:
            ws.send_json({
                "type": "message",
                "content": "Hello, world!",
                "model": "llama3",
                "backend": "ollama",
            })

            # Should receive agent_joined
            msg1 = ws.receive_json()
            assert msg1["type"] == "agent_joined"
            assert msg1["agent"] == "llama3"

            # Should receive chunk
            msg2 = ws.receive_json()
            assert msg2["type"] == "chunk"
            assert msg2["content"] == "Mock LLM response"
            assert msg2["model"] == "llama3"

            # Should receive complete
            msg3 = ws.receive_json()
            assert msg3["type"] == "complete"
            assert msg3["content"] == "Mock LLM response"
            assert msg3["model"] == "llama3"
            assert "elapsed_ms" in msg3

    def test_ws_empty_message_error(self, test_client):
        with test_client.websocket_connect("/ws/chat") as ws:
            ws.send_json({
                "type": "message",
                "content": "",
            })
            msg = ws.receive_json()
            assert msg["type"] == "error"
            assert "Empty" in msg["message"]

    def test_ws_ping_pong(self, test_client):
        with test_client.websocket_connect("/ws/chat") as ws:
            ws.send_json({"type": "ping"})
            msg = ws.receive_json()
            assert msg["type"] == "pong"

    def test_ws_auto_select_model(self, test_client):
        with test_client.websocket_connect("/ws/chat") as ws:
            ws.send_json({
                "type": "message",
                "content": "Hello",
            })

            msg1 = ws.receive_json()
            assert msg1["type"] == "agent_joined"
            # Should auto-select llama3
            assert msg1["agent"] == "llama3"

            # Consume chunk and complete
            ws.receive_json()
            ws.receive_json()

    def test_ws_llm_error(self, test_client, mock_llm_router):
        mock_llm_router.side_effect = RuntimeError("Model crashed")
        with test_client.websocket_connect("/ws/chat") as ws:
            ws.send_json({
                "type": "message",
                "content": "Hello",
                "model": "llama3",
                "backend": "ollama",
            })

            # agent_joined still sent before the call
            msg1 = ws.receive_json()
            assert msg1["type"] == "agent_joined"

            # Then error
            msg2 = ws.receive_json()
            assert msg2["type"] == "error"
            assert "Model crashed" in msg2["message"]

    def test_ws_multiple_messages(self, test_client):
        with test_client.websocket_connect("/ws/chat") as ws:
            for i in range(3):
                ws.send_json({
                    "type": "message",
                    "content": f"Message {i}",
                    "model": "llama3",
                    "backend": "ollama",
                })
                # Consume agent_joined + chunk + complete
                ws.receive_json()  # agent_joined
                ws.receive_json()  # chunk
                ws.receive_json()  # complete


# ---------------------------------------------------------------------------
# Helper function tests
# ---------------------------------------------------------------------------

class TestHelperFunctions:
    def test_build_messages_basic(self):
        from brynq.api_server import _build_messages
        msgs = _build_messages("Hello")
        assert len(msgs) == 1
        assert msgs[0]["role"] == "user"
        assert msgs[0]["content"] == "Hello"

    def test_build_messages_with_system(self):
        from brynq.api_server import _build_messages
        msgs = _build_messages("Hello", system_prompt="Be helpful")
        assert len(msgs) == 2
        assert msgs[0]["role"] == "system"
        assert msgs[0]["content"] == "Be helpful"
        assert msgs[1]["role"] == "user"

    def test_build_messages_with_history(self):
        from brynq.api_server import _build_messages
        history = [
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": "Hello!"},
        ]
        msgs = _build_messages("Follow up", conversation=history)
        assert len(msgs) == 3
        assert msgs[0]["role"] == "user"
        assert msgs[0]["content"] == "Hi"
        assert msgs[1]["role"] == "assistant"
        assert msgs[2]["role"] == "user"
        assert msgs[2]["content"] == "Follow up"

    def test_build_messages_filters_empty(self):
        from brynq.api_server import _build_messages
        history = [
            {"role": "user", "content": "Hi"},
            {"role": "user", "content": ""},  # empty, should be skipped
        ]
        msgs = _build_messages("Now", conversation=history)
        assert len(msgs) == 2  # "Hi" + "Now"

    def test_resolve_backend_model_explicit(self):
        from brynq.api_server import _resolve_backend_model
        with patch("brynq.llm_router._default_model_for", return_value="llama3"):
            b, m = _resolve_backend_model("ollama", "mistral")
            assert b == "ollama"
            assert m == "mistral"

    def test_resolve_backend_model_auto(self):
        from brynq.api_server import _resolve_backend_model
        with patch("brynq.llm_router._pick_first_available", return_value=("ollama", "llama3")):
            b, m = _resolve_backend_model(None, None)
            assert b == "ollama"
            assert m == "llama3"


# ---------------------------------------------------------------------------
# OpenAI schema registration
# ---------------------------------------------------------------------------

class TestChatSchemaRegistration:
    def test_chat_endpoints_registered(self):
        from brynq.api_server import _OPENAI_ENDPOINT_REGISTRY
        paths = [ep["path"] for ep in _OPENAI_ENDPOINT_REGISTRY]
        assert "/api/chat" in paths
        assert "/api/chat/chain" in paths
        assert "/api/models" in paths
        assert "/api/ollama/status" in paths

    def test_chat_group_exists(self):
        from brynq.api_server import _OPENAI_ENDPOINT_REGISTRY
        groups = set(ep.get("group", "") for ep in _OPENAI_ENDPOINT_REGISTRY)
        assert "chat" in groups
