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

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

# ==========================================
# 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


@patch("dracula.client.genai.Client")
@patch("pathlib.Path.read_bytes")
@patch("pathlib.Path.exists")
def test_dracula_multimodal_chat(mock_exists, mock_read, MockClient, tmp_path):
    """Test if chat correctly handles multimodal file inputs with mocks."""
    mock_exists.return_value = True
    mock_read.return_value = b"fake_image_data"

    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("I see a cat.")

    ai = Dracula(api_key="fake-key", db_path=tmp_path / "multimodal.db")
    reply = ai.chat("What is this?", filepath="cat.jpg")

    assert reply == "I see a cat."
    history = ai.get_history()
    assert "[Attached: cat.jpg]" in history[-2]["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


@pytest.mark.asyncio
@patch("dracula.async_client.genai.Client")
@patch("pathlib.Path.read_bytes")
@patch("pathlib.Path.exists")
async def test_async_dracula_multimodal_chat(
    mock_exists, mock_read, MockClient, tmp_path
):
    """Test if async chat correctly handles multimodal file inputs."""
    mock_exists.return_value = True
    mock_read.return_value = b"fake_async_data"

    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 vision works!")

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

    reply = await ai.chat("Look at this", filepath="test.png")
    assert reply == "Async vision works!"

    history = await ai.get_history()
    assert "[Attached: test.png]" in history[-2]["content"]


# ==========================================
# validate_model tests
# ==========================================


def test_validate_model_allows_custom_string():
    """Any non-empty string should be accepted as a model name."""
    from dracula.validators import validate_model

    validate_model("gemini-custom-preview")
    validate_model("gemini-2.5-flash-exp-0325")


def test_validate_model_rejects_empty_string():
    """An empty or whitespace-only string must raise ValidationException."""
    from dracula.validators import validate_model

    with pytest.raises(ValidationException):
        validate_model("")

    with pytest.raises(ValidationException):
        validate_model("   ")


def test_validate_model_rejects_non_string():
    """Non-string, non-enum values must raise ValidationException."""
    from dracula.validators import validate_model

    with pytest.raises(ValidationException):
        validate_model(42)


# ==========================================
# session_id isolation via Dracula client
# ==========================================


@patch("dracula.client.genai.Client")
def test_dracula_session_id_isolation(MockClient, tmp_path):
    """Two Dracula instances with the same db_path but different session_ids
    must have fully isolated conversation histories."""
    db_file = tmp_path / "shared_sessions.db"

    ai_a = Dracula(api_key="fake-key", db_path=db_file, session_id="user-001")
    ai_b = Dracula(api_key="fake-key", db_path=db_file, session_id="user-002")

    ai_a.memory.add_message("user", "Message for user 001")
    ai_b.memory.add_message("user", "Message for user 002")

    history_a = ai_a.memory.get_history()
    history_b = ai_b.memory.get_history()

    assert len(history_a) == 1
    assert history_a[0]["content"] == "Message for user 001"
    assert len(history_b) == 1
    assert history_b[0]["content"] == "Message for user 002"


# ==========================================
# Pricing / estimate_cost unit tests
# ==========================================


def test_estimate_cost_known_model():
    """estimate_cost() should compute the right USD amount for a known model."""
    # gemini-2.5-flash: input=$0.075/M, output=$0.30/M
    cost = estimate_cost("gemini-2.5-flash", 1_000_000, 1_000_000)
    assert cost == round((0.075 + 0.30), 6)


def test_estimate_cost_unknown_model_uses_default():
    """estimate_cost() should fall back to default pricing for unknown models."""
    # default: input=$0.10/M, output=$0.40/M
    cost = estimate_cost("unknown-model-xyz", 1_000_000, 0)
    assert cost == 0.10


def test_estimate_cost_zero_tokens():
    """estimate_cost() with zero tokens should return 0.0."""
    assert estimate_cost("gemini-2.5-flash", 0, 0) == 0.0


# ==========================================
# Before/After hook tests
# ==========================================


@patch("dracula.client.genai.Client")
def test_before_chat_hook_transforms_message(MockClient, tmp_path):
    """A before_chat hook that returns a new string should replace the message."""
    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("HOOK REPLY")

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

    received = []

    @ai.before_chat
    def capture(msg):
        received.append(msg)
        return msg.upper()

    mock_session.send_message.return_value = create_mock_response("reply")
    ai.chat("hello")

    # The hook should have received the original and uppercased it.
    assert received == ["hello"]
    # send_message was called — check it received the uppercased version.
    call_args = mock_session.send_message.call_args[0][0]
    assert call_args == "HELLO"


@patch("dracula.client.genai.Client")
def test_after_chat_hook_transforms_reply(MockClient, tmp_path):
    """An after_chat hook that returns a new string should replace the response."""
    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("original reply")

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

    @ai.after_chat
    def shout(msg, reply):
        return reply.upper()

    result = ai.chat("hi")
    assert result == "ORIGINAL REPLY"


@patch("dracula.client.genai.Client")
def test_after_chat_hook_returning_none_keeps_original(MockClient, tmp_path):
    """An after_chat hook that returns None should leave the reply unchanged."""
    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("keep me")

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

    @ai.after_chat
    def noop(msg, reply):
        pass  # returns None implicitly

    result = ai.chat("hi")
    assert result == "keep me"


# ==========================================
# Token budget tests
# ==========================================


@patch("dracula.client.genai.Client")
def test_token_budget_exceeded_raises(MockClient, tmp_path):
    """chat() should raise BudgetExceededException when cumulative tokens >= budget."""
    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("reply")

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

    # Manually record tokens to simulate budget exhaustion.
    ai.stats.record_tokens(60, 50)  # total = 110 >= 100

    with pytest.raises(BudgetExceededException):
        ai.chat("this should fail")


@patch("dracula.client.genai.Client")
def test_token_budget_not_exceeded_allows_chat(MockClient, tmp_path):
    """chat() should proceed when cumulative tokens are still under budget."""
    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("allowed reply")

    ai = Dracula(api_key="fake-key", db_path=tmp_path / "mem.db", token_budget=1000)
    ai.stats.record_tokens(10, 10)  # total = 20 < 1000

    result = ai.chat("this should work")
    assert result == "allowed reply"


# ==========================================
# Response cache via Dracula.chat()
# ==========================================


@patch("dracula.client.genai.Client")
def test_response_cache_hit_skips_api(MockClient, tmp_path):
    """A cache hit should return the cached reply without calling send_message."""
    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("api reply")

    cache_db = tmp_path / "cache.db"
    ai = Dracula(
        api_key="fake-key",
        db_path=tmp_path / "mem.db",
        cache_ttl=60,
        cache_db_path=cache_db,
    )

    # Prime the cache manually.
    ai._cache.set(ai.memory.session_id, "cached question", "cached answer")

    result = ai.chat("cached question")

    assert result == "cached answer"
    mock_session.send_message.assert_not_called()


# ==========================================
# estimated_cost()
# ==========================================


@patch("dracula.client.genai.Client")
def test_estimated_cost_returns_float(MockClient, tmp_path):
    """estimated_cost() should return a non-negative float."""
    ai = Dracula(api_key="fake-key", db_path=tmp_path / "mem.db")
    ai.stats.record_tokens(500, 1000)

    cost = ai.estimated_cost()
    assert isinstance(cost, float)
    assert cost >= 0.0


@patch("dracula.client.genai.Client")
def test_estimated_cost_zero_when_no_tokens(MockClient, tmp_path):
    """estimated_cost() should be 0.0 when no tokens have been used."""
    ai = Dracula(api_key="fake-key", db_path=tmp_path / "mem.db", stats_db_path=tmp_path / "stats.db")
    assert ai.estimated_cost() == 0.0


# ==========================================
# Fork tests
# ==========================================


@patch("dracula.client.genai.Client")
def test_fork_copies_history(MockClient, tmp_path):
    """fork() should create a new instance with the same history."""
    ai = Dracula(
        api_key="fake-key",
        db_path=tmp_path / "mem.db",
        session_id="original",
    )
    ai.memory.add_message("user", "Hello")
    ai.memory.add_message("model", "Hi there!")

    branch = ai.fork(session_id="branch-1")

    branch_history = branch.get_history()
    assert len(branch_history) == 2
    assert branch_history[0]["content"] == "Hello"
    assert branch_history[1]["content"] == "Hi there!"


@patch("dracula.client.genai.Client")
def test_fork_is_isolated_from_original(MockClient, tmp_path):
    """Changes to the fork's history must not affect the original."""
    ai = Dracula(
        api_key="fake-key",
        db_path=tmp_path / "mem.db",
        session_id="original",
    )
    ai.memory.add_message("user", "Shared message")

    branch = ai.fork(session_id="branch-isolated")
    branch.memory.add_message("user", "Branch-only message")

    original_history = ai.get_history()
    branch_history = branch.get_history()

    assert len(original_history) == 1
    assert len(branch_history) == 2


# ==========================================
# Async hook tests
# ==========================================


@pytest.mark.asyncio
@patch("dracula.async_client.genai.Client")
async def test_async_before_hook(MockClient, tmp_path):
    """AsyncDracula before_chat hook should transform the message."""
    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 reply")

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

    received = []

    @ai.before_chat
    def capture(msg):
        received.append(msg)
        return "TRANSFORMED"

    await ai.chat("original")
    assert received == ["original"]
    call_args = mock_session.send_message.call_args[0][0]
    assert call_args == "TRANSFORMED"


@pytest.mark.asyncio
@patch("dracula.async_client.genai.Client")
async def test_async_after_hook(MockClient, tmp_path):
    """AsyncDracula after_chat hook should transform the reply."""
    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("raw reply")

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

    @ai.after_chat
    def shout(msg, reply):
        return reply.upper()

    result = await ai.chat("hi")
    assert result == "RAW REPLY"


# ==========================================
# Async estimated_cost
# ==========================================


@pytest.mark.asyncio
@patch("dracula.async_client.genai.Client")
async def test_async_estimated_cost(MockClient, tmp_path):
    """AsyncDracula.estimated_cost() should return a non-negative float."""
    ai = AsyncDracula(api_key="fake-key", db_path=tmp_path / "async_mem.db")
    await ai._ensure_initialized()
    await ai.stats.record_tokens(200, 400)

    cost = await ai.estimated_cost()
    assert isinstance(cost, float)
    assert cost >= 0.0


# ==========================================
# Async fork test
# ==========================================


@pytest.mark.asyncio
@patch("dracula.async_client.genai.Client")
async def test_async_fork_copies_history(MockClient, tmp_path):
    """AsyncDracula.fork() should create a branch with the same history."""
    ai = AsyncDracula(
        api_key="fake-key",
        db_path=tmp_path / "async_mem.db",
        session_id="async-original",
    )
    await ai._ensure_initialized()
    await ai.memory.add_message("user", "Async Hello")
    await ai.memory.add_message("model", "Async Hi!")

    branch = await ai.fork(session_id="async-branch")

    branch_history = await branch.get_history()
    assert len(branch_history) == 2
    assert branch_history[0]["content"] == "Async Hello"
