import pytest
import asyncio
from unittest.mock import patch, MagicMock, AsyncMock

from dracula import Dracula, AsyncDracula
from dracula.exceptions import (
    InvalidAPIKeyException,
    PersonaException,
    ValidationException,
)

# ==========================================
# Helpers for Mocking Gemini API Responses
# ==========================================


def create_mock_response(text_content="Mocked response"):
    """Creates a mock Gemini API response to bypass safety checks."""
    mock_part = MagicMock()
    mock_part.function_call = None

    mock_content = MagicMock()
    mock_content.parts = [mock_part]

    mock_candidate = MagicMock()
    mock_candidate.content = mock_content

    mock_response = MagicMock()
    mock_response.text = text_content
    mock_response.candidates = [mock_candidate]

    return mock_response


# ==========================================
# 1. Tests for synchronous Dracula class
# ==========================================


@patch("dracula.client.genai.Client")
def test_dracula_initialization(MockClient, tmp_path):
    """Test Dracula initialization and API key validation."""
    db_file = tmp_path / "mem.db"

    ai = Dracula(api_key="fake-key", db_path=db_file)
    assert ai.prompt == "You are a helpful assistant."
    assert ai.language == "English"

    with pytest.raises(ValidationException):
        Dracula(api_key="")


@patch("dracula.client.genai.Client")
def test_dracula_chat(MockClient, tmp_path):
    """Test standard chat functionality and stat recording."""
    mock_client_instance = MockClient.return_value
    mock_session = MagicMock()
    mock_client_instance.chats.create.return_value = mock_session
    mock_session.send_message.return_value = create_mock_response("Hello from Gemini")

    ai = Dracula(api_key="fake-key", db_path=tmp_path / "mem.db")
    ai.reset_stats()

    reply = ai.chat("Hi")
    assert reply == "Hello from Gemini"

    stats = ai.get_stats()
    assert stats["total_messages"] == 1
    assert stats["total_responses"] == 1


@patch("dracula.client.genai.Client")
def test_dracula_stream(MockClient, tmp_path):
    """Test streaming chat functionality."""
    mock_client_instance = MockClient.return_value
    mock_session = MagicMock()
    mock_client_instance.chats.create.return_value = mock_session
    mock_session.send_message_stream.return_value = [
        MagicMock(text="Chunk 1 "),
        MagicMock(text="Chunk 2"),
    ]

    ai = Dracula(api_key="fake-key", db_path=tmp_path / "mem.db")

    chunks = list(ai.stream("Tell me a story"))
    assert len(chunks) == 2
    assert chunks[0] == "Chunk 1 "

    history = ai.get_history()
    assert history[-1]["content"] == "Chunk 1 Chunk 2"


def test_dracula_set_persona(tmp_path):
    """Test if switching personas correctly updates attributes and clears memory."""
    ai = Dracula(api_key="fake-key", db_path=tmp_path / "mem.db")
    ai.set_persona("pirate")
    assert "pirate" in ai.prompt.lower() or "arr" in ai.prompt.lower()

    with pytest.raises(PersonaException):
        ai.set_persona("unknown_persona_123")


@patch("dracula.client.genai.Client")
def test_dracula_count_tokens(MockClient, tmp_path):
    """Test token counting functionality."""
    mock_client_instance = MockClient.return_value

    mock_response = MagicMock()
    mock_response.total_tokens = 42
    mock_client_instance.models.count_tokens.return_value = mock_response

    ai = Dracula(api_key="fake-key", db_path=tmp_path / "mem.db")

    tokens = ai.count_tokens("Hello world")
    assert tokens == 42

    ai.memory.add_message("user", "Hello history")
    total_tokens = ai.count_history_tokens()
    assert total_tokens == 42


@patch("dracula.client.genai.Client")
def test_dracula_export_markdown(MockClient, tmp_path):
    """Test exporting conversation history to a Markdown file."""
    db_file = tmp_path / "mem.db"
    export_file = tmp_path / "chat.md"

    ai = Dracula(api_key="fake-key", db_path=db_file)

    ai.memory.add_message("user", "Hello Markdown!")
    ai.memory.add_message("model", "This is a bold **response**.")

    ai.export_to_markdown(str(export_file))

    assert export_file.exists()
    content = export_file.read_text(encoding="utf-8")

    assert "# Dracula Conversation History\n\n" in content
    assert "👤 **You**" in content
    assert "Hello Markdown!" in content
    assert "🤖 **Gemini**" in content
    assert "This is a bold **response**." in content


# ==========================================
# 2. Tests for AsyncDracula class
# ==========================================


@pytest.mark.asyncio
@patch("dracula.async_client.genai.Client")
async def test_async_dracula_chat(MockClient, tmp_path):
    """Test async chat functionality."""
    mock_client_instance = MockClient.return_value
    mock_session = AsyncMock()
    mock_client_instance.aio.chats.create.return_value = mock_session
    mock_session.send_message.return_value = create_mock_response("Async Hello!")

    ai = AsyncDracula(api_key="fake-key", db_path=tmp_path / "async_mem.db")
    await ai._ensure_initialized()
    await ai.reset_stats()

    reply = await ai.chat("Hi async")
    assert reply == "Async Hello!"

    stats = await ai.get_stats()
    assert stats["total_messages"] == 1

    history = await ai.get_history()
    assert len(history) == 2


@pytest.mark.asyncio
@patch("dracula.async_client.genai.Client")
async def test_async_dracula_stream(MockClient, tmp_path):
    """Test async streaming chat functionality."""
    mock_client_instance = MockClient.return_value
    mock_session = AsyncMock()
    mock_client_instance.aio.chats.create.return_value = mock_session

    async def mock_async_gen():
        yield MagicMock(text="Async ")
        yield MagicMock(text="Chunk")

    mock_session.send_message_stream.return_value = mock_async_gen()

    ai = AsyncDracula(api_key="fake-key", db_path=tmp_path / "async_mem.db")
    await ai._ensure_initialized()

    chunks = []
    async for chunk in ai.stream("Async story"):
        chunks.append(chunk)

    assert len(chunks) == 2

    history = await ai.get_history()
    assert history[-1]["content"] == "Async Chunk"


@pytest.mark.asyncio
@patch("dracula.async_client.genai.Client")
async def test_async_context_manager(MockClient, tmp_path):
    """Test if async context manager cleans up properly."""
    db_file = tmp_path / "async_ctx.db"

    async with AsyncDracula(api_key="fake-key", db_path=db_file) as ai:
        await ai._ensure_initialized()
        await ai.memory.add_message("user", "test")
        history = await ai.get_history()
        assert len(history) == 1

    from dracula.memory import AsyncMemory

    check_memory = AsyncMemory(db_path=db_file)
    await check_memory.initialize()
    history_after = await check_memory.get_history()
    assert len(history_after) == 0


@pytest.mark.asyncio
@patch("dracula.async_client.genai.Client")
async def test_async_dracula_count_tokens(MockClient, tmp_path):
    """Test async token counting functionality."""
    mock_client_instance = MockClient.return_value

    mock_response = MagicMock()
    mock_response.total_tokens = 84

    mock_client_instance.aio.models.count_tokens = AsyncMock(return_value=mock_response)

    ai = AsyncDracula(api_key="fake-key", db_path=tmp_path / "async_mem.db")
    await ai._ensure_initialized()

    tokens = await ai.count_tokens("Async hello")
    assert tokens == 84

    await ai.memory.add_message("user", "Async history test")
    total_tokens = await ai.count_history_tokens()
    assert total_tokens == 84


@pytest.mark.asyncio
@patch("dracula.async_client.genai.Client")
async def test_async_dracula_export_markdown(MockClient, tmp_path):
    """Test asynchronously exporting conversation history to a Markdown file."""
    db_file = tmp_path / "async_mem.db"
    export_file = tmp_path / "async_chat.md"

    ai = AsyncDracula(api_key="fake-key", db_path=db_file)
    await ai._ensure_initialized()

    await ai.memory.add_message("user", "Async Markdown Test")
    await ai.memory.add_message("model", "Async response here.")

    await ai.export_to_markdown(str(export_file))

    assert export_file.exists()
    content = export_file.read_text(encoding="utf-8")

    assert "# Dracula Conversation History\n\n" in content
    assert "👤 **You**" in content
    assert "Async Markdown Test" in content
    assert "🤖 **Gemini**" in content
