import pytest
import json
import sqlite3
import time
from pathlib import Path

from dracula.memory import Memory, AsyncMemory
from dracula.stats import Stats, AsyncStats
from dracula.cache import Cache, AsyncCache
from dracula.exceptions import ChatException

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


def test_memory_add_and_get_history(tmp_path):
    """Test adding messages and retrieving them chronologically."""
    db_file = tmp_path / "memory.db"
    memory = Memory(db_path=db_file, max_messages=5)

    memory.add_message("user", "Hello")
    memory.add_message("model", "Hi there!")

    history = memory.get_history()
    assert len(history) == 2
    assert history[0] == {"role": "user", "content": "Hello"}
    assert history[1] == {"role": "model", "content": "Hi there!"}


def test_memory_trim_history(tmp_path):
    """Test that history is trimmed correctly when exceeding max_messages."""
    db_file = tmp_path / "memory.db"
    memory = Memory(db_path=db_file, max_messages=3)

    for i in range(5):
        memory.add_message("user", f"Message {i}")

    history = memory.get_history()

    assert len(history) == 3
    assert history[0]["content"] == "Message 2"
    assert history[-1]["content"] == "Message 4"


def test_memory_clear(tmp_path):
    """Test permanently deleting the conversation history."""
    db_file = tmp_path / "memory.db"
    memory = Memory(db_path=db_file)

    memory.add_message("user", "Hello")
    assert len(memory.get_history()) == 1

    memory.clear()
    assert len(memory.get_history()) == 0


def test_memory_save_and_load(tmp_path):
    """Test exporting to JSON and migrating back to SQLite."""
    db_file = tmp_path / "memory.db"
    json_file = tmp_path / "backup.json"
    memory = Memory(db_path=db_file)

    memory.add_message("user", "Test message")
    memory.save(filepath=str(json_file))

    assert json_file.exists()

    db_file2 = tmp_path / "memory2.db"
    memory2 = Memory(db_path=db_file2)
    memory2.load(filepath=str(json_file))

    history = memory2.get_history()
    assert len(history) == 1
    assert history[0]["content"] == "Test message"


# ==========================================
# 2. Tests for synchronous Stats class
# ==========================================


def test_stats_initialization(tmp_path):
    """Test if stats are initialized to zero."""
    db_file = tmp_path / "stats.db"
    stats = Stats(db_path=db_file)

    data = stats.get_stats()
    assert data["total_messages"] == 0
    assert data["total_characters_sent"] == 0


def test_stats_recording_and_reset(tmp_path):
    """Test recording messages/responses and resetting."""
    db_file = tmp_path / "stats.db"
    stats = Stats(db_path=db_file)

    stats.record_message("Hello")
    stats.record_response("Hi there!")

    data = stats.get_stats()
    assert data["total_messages"] == 1
    assert data["total_characters_sent"] == 5
    assert data["total_responses"] == 1
    assert data["total_characters_received"] == 9

    stats.reset()
    reset_data = stats.get_stats()
    assert reset_data["total_messages"] == 0
    assert reset_data["total_responses"] == 0


# ==========================================
# 3. Tests for AsyncMemory class
# ==========================================


@pytest.mark.asyncio
async def test_async_memory_add_get_clear(tmp_path):
    """Test async operations for conversation history."""
    db_file = tmp_path / "async_memory.db"
    memory = AsyncMemory(db_path=db_file, max_messages=2)
    await memory.initialize()

    await memory.add_message("user", "Async Hello")
    await memory.add_message("model", "Async Hi")
    await memory.add_message("user", "Overflow message")

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

    assert history[0]["content"] == "Async Hi"

    await memory.clear()
    empty_history = await memory.get_history()
    assert len(empty_history) == 0


# ==========================================
# 4. Tests for AsyncStats class
# ==========================================


@pytest.mark.asyncio
async def test_async_stats_record_and_reset(tmp_path):
    """Test async operations for usage tracking."""
    db_file = tmp_path / "async_stats.db"
    stats = AsyncStats(db_path=db_file)
    await stats.initialize()

    await stats.record_message("Hi")
    await stats.record_response("Hello!")

    data = await stats.get_stats()
    assert data["total_messages"] == 1
    assert data["total_characters_sent"] == 2
    assert data["total_responses"] == 1
    assert data["total_characters_received"] == 6

    await stats.reset()
    reset_data = await stats.get_stats()
    assert reset_data["total_messages"] == 0
    assert reset_data["total_characters_received"] == 0


# ==========================================
# 5. Session isolation tests
# ==========================================


def test_memory_session_isolation(tmp_path):
    """Two Memory instances sharing a db_path must not bleed into each other."""
    db_file = tmp_path / "shared.db"
    session_a = Memory(db_path=db_file, session_id="session-a")
    session_b = Memory(db_path=db_file, session_id="session-b")

    session_a.add_message("user", "Message from A")
    session_b.add_message("user", "Message from B")

    history_a = session_a.get_history()
    history_b = session_b.get_history()

    assert len(history_a) == 1
    assert history_a[0]["content"] == "Message from A"
    assert len(history_b) == 1
    assert history_b[0]["content"] == "Message from B"


@pytest.mark.asyncio
async def test_async_memory_session_isolation(tmp_path):
    """Two AsyncMemory instances sharing a db_path must not bleed into each other."""
    db_file = tmp_path / "shared_async.db"
    session_a = AsyncMemory(db_path=db_file, session_id="async-session-a")
    session_b = AsyncMemory(db_path=db_file, session_id="async-session-b")
    await session_a.initialize()
    await session_b.initialize()

    await session_a.add_message("user", "Async message from A")
    await session_b.add_message("user", "Async message from B")

    history_a = await session_a.get_history()
    history_b = await session_b.get_history()

    assert len(history_a) == 1
    assert history_a[0]["content"] == "Async message from A"
    assert len(history_b) == 1
    assert history_b[0]["content"] == "Async message from B"


def test_memory_migration(tmp_path):
    """Memory must migrate an old database that lacks the session_id column."""
    db_file = tmp_path / "old.db"

    # Create a legacy database without the session_id column.
    with sqlite3.connect(db_file) as conn:
        conn.execute(
            """
            CREATE TABLE history (
                id        INTEGER PRIMARY KEY AUTOINCREMENT,
                role      TEXT NOT NULL,
                content   TEXT NOT NULL,
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
            )
            """
        )
        conn.execute("INSERT INTO history (role, content) VALUES ('user', 'legacy msg')")
        conn.commit()

    # Initialising Memory should apply the migration without raising.
    memory = Memory(db_path=db_file, session_id="migrated-session")

    # The migrated column must be present.
    with sqlite3.connect(db_file) as conn:
        cols = [row[1] for row in conn.execute("PRAGMA table_info(history)")]
    assert "session_id" in cols

    # New messages should be scoped correctly.
    memory.add_message("user", "post-migration msg")
    history = memory.get_history()
    assert any(m["content"] == "post-migration msg" for m in history)


# ==========================================
# 6. Tests for Cache (sync)
# ==========================================


def test_cache_hit(tmp_path):
    """A stored response should be returned on a matching get()."""
    cache = Cache(db_path=tmp_path / "cache.db", ttl=60)
    cache.set("session-1", "hello world", "cached reply")
    result = cache.get("session-1", "hello world")
    assert result == "cached reply"


def test_cache_miss_different_session(tmp_path):
    """A cached entry for session-A must not be returned for session-B."""
    cache = Cache(db_path=tmp_path / "cache.db", ttl=60)
    cache.set("session-A", "same question", "reply A")
    result = cache.get("session-B", "same question")
    assert result is None


def test_cache_miss_unknown_message(tmp_path):
    """A get() for a message that was never cached returns None."""
    cache = Cache(db_path=tmp_path / "cache.db", ttl=60)
    result = cache.get("session-1", "never stored")
    assert result is None


def test_cache_disabled_always_returns_none(tmp_path):
    """Cache with ttl=0 must always return None (caching disabled)."""
    cache = Cache(db_path=tmp_path / "cache.db", ttl=0)
    cache.set("session-1", "hello", "should not be stored")
    result = cache.get("session-1", "hello")
    assert result is None


def test_cache_invalidate(tmp_path):
    """invalidate() removes all entries for the given session."""
    cache = Cache(db_path=tmp_path / "cache.db", ttl=60)
    cache.set("session-1", "q1", "r1")
    cache.set("session-1", "q2", "r2")
    cache.set("session-2", "q1", "r1-other")

    cache.invalidate("session-1")

    assert cache.get("session-1", "q1") is None
    assert cache.get("session-1", "q2") is None
    # Session-2 must be untouched.
    assert cache.get("session-2", "q1") == "r1-other"


def test_cache_expired_entry_returns_none(tmp_path):
    """An entry whose TTL has elapsed must not be returned."""
    cache = Cache(db_path=tmp_path / "cache.db", ttl=1)
    cache.set("session-1", "expiring question", "expiring reply")
    time.sleep(1.1)
    result = cache.get("session-1", "expiring question")
    assert result is None


# ==========================================
# 7. Tests for AsyncCache
# ==========================================


@pytest.mark.asyncio
async def test_async_cache_hit(tmp_path):
    """Async cache: a stored response should be returned on a matching get()."""
    cache = AsyncCache(db_path=tmp_path / "async_cache.db", ttl=60)
    await cache.initialize()
    await cache.set("session-1", "async hello", "async cached reply")
    result = await cache.get("session-1", "async hello")
    assert result == "async cached reply"


@pytest.mark.asyncio
async def test_async_cache_disabled(tmp_path):
    """Async cache with ttl=0 must always return None."""
    cache = AsyncCache(db_path=tmp_path / "async_cache.db", ttl=0)
    await cache.initialize()
    await cache.set("session-1", "hello", "should not store")
    result = await cache.get("session-1", "hello")
    assert result is None


@pytest.mark.asyncio
async def test_async_cache_invalidate(tmp_path):
    """Async invalidate() removes only that session's entries."""
    cache = AsyncCache(db_path=tmp_path / "async_cache.db", ttl=60)
    await cache.initialize()
    await cache.set("sess-A", "q", "r")
    await cache.set("sess-B", "q", "r-other")
    await cache.invalidate("sess-A")
    assert await cache.get("sess-A", "q") is None
    assert await cache.get("sess-B", "q") == "r-other"


# ==========================================
# 8. Tests for Stats token tracking
# ==========================================


def test_stats_record_tokens(tmp_path):
    """record_tokens() should accumulate in the DB and be returned by get_stats()."""
    stats = Stats(db_path=tmp_path / "stats.db")
    stats.record_tokens(100, 200)
    stats.record_tokens(50, 75)

    data = stats.get_stats()
    assert data["total_input_tokens"] == 150
    assert data["total_output_tokens"] == 275


def test_stats_reset_clears_tokens(tmp_path):
    """reset() must zero token columns alongside message counts."""
    stats = Stats(db_path=tmp_path / "stats.db")
    stats.record_tokens(500, 1000)
    stats.reset()

    data = stats.get_stats()
    assert data["total_input_tokens"] == 0
    assert data["total_output_tokens"] == 0


@pytest.mark.asyncio
async def test_async_stats_record_tokens(tmp_path):
    """AsyncStats.record_tokens() should accumulate correctly."""
    stats = AsyncStats(db_path=tmp_path / "async_stats.db")
    await stats.initialize()
    await stats.record_tokens(300, 600)
    await stats.record_tokens(100, 200)

    data = await stats.get_stats()
    assert data["total_input_tokens"] == 400
    assert data["total_output_tokens"] == 800


@pytest.mark.asyncio
async def test_async_stats_reset_clears_tokens(tmp_path):
    """AsyncStats.reset() must zero token columns."""
    stats = AsyncStats(db_path=tmp_path / "async_stats.db")
    await stats.initialize()
    await stats.record_tokens(999, 999)
    await stats.reset()

    data = await stats.get_stats()
    assert data["total_input_tokens"] == 0
    assert data["total_output_tokens"] == 0


# ==========================================
# 9. Tests for Memory.compress_to_summary
# ==========================================


def test_memory_compress_to_summary(tmp_path):
    """compress_to_summary() replaces old messages with one summary entry."""
    db_file = tmp_path / "compress.db"
    memory = Memory(db_path=db_file, max_messages=20)

    for i in range(6):
        memory.add_message("user", f"Message {i}")
        memory.add_message("model", f"Reply {i}")

    # Compress keeping last 2 messages.
    memory.compress_to_summary("This is a summary.", keep_last_n=2)

    history = memory.get_history()
    # Should have: 1 summary + 2 kept = 3 entries.
    assert len(history) == 3
    assert history[0]["role"] == "user"
    assert "[Context Summary]" in history[0]["content"]
    assert "This is a summary." in history[0]["content"]
    # Last two messages should be the most recent ones.
    assert history[1]["content"] == "Message 5"
    assert history[2]["content"] == "Reply 5"


@pytest.mark.asyncio
async def test_async_memory_compress_to_summary(tmp_path):
    """AsyncMemory.compress_to_summary() replaces old messages with a summary."""
    db_file = tmp_path / "async_compress.db"
    memory = AsyncMemory(db_path=db_file, max_messages=20)
    await memory.initialize()

    for i in range(4):
        await memory.add_message("user", f"Async msg {i}")
        await memory.add_message("model", f"Async reply {i}")

    await memory.compress_to_summary("Async summary here.", keep_last_n=2)

    history = await memory.get_history()
    assert len(history) == 3
    assert "[Context Summary]" in history[0]["content"]
    assert "Async summary here." in history[0]["content"]
