import uuid
import os
import memgpt.utils as utils

utils.DEBUG = True
from memgpt.config import MemGPTConfig
from memgpt.server.server import SyncServer
from memgpt.data_types import EmbeddingConfig, AgentState, LLMConfig, Message, Passage
from memgpt.embeddings import embedding_model
from memgpt.metadata import MetadataStore
from .utils import wipe_config, wipe_memgpt_home


def test_server():
    wipe_memgpt_home()

    config = MemGPTConfig.load()

    # setup config for postgres storage
    config.archival_storage_uri = os.getenv("PGVECTOR_TEST_DB_URL")
    config.recall_storage_uri = os.getenv("PGVECTOR_TEST_DB_URL")
    config.archival_storage_type = "postgres"
    config.recall_storage_type = "postgres"
    config.save()

    user_id = uuid.UUID(config.anon_clientid)
    ms = MetadataStore(config)
    server = SyncServer()

    try:
        fake_agent_id = uuid.uuid4()
        server.user_message(user_id=user_id, agent_id=fake_agent_id, message="Hello?")
        raise Exception("user_message call should have failed")
    except (KeyError, ValueError) as e:
        # Error is expected
        print(e)
    except:
        raise

    # embedding config
    if os.getenv("OPENAI_API_KEY"):
        embedding_config = EmbeddingConfig(
            embedding_endpoint_type="openai",
            embedding_endpoint="https://api.openai.com/v1",
            embedding_dim=1536,
            openai_key=os.getenv("OPENAI_API_KEY"),
        )
        print("Using OpenAI embeddings")
    else:
        embedding_config = EmbeddingConfig(embedding_endpoint_type="local", embedding_endpoint=None, embedding_dim=384)
        print("Using local embeddings")

    agent_state = server.create_agent(
        user_id=user_id,
        agent_config=dict(
            name="test_agent", user_id=user_id, preset="memgpt_chat", human="cs_phd", persona="sam_pov", embedding_config=embedding_config
        ),
    )
    print(f"Created agent\n{agent_state}")

    try:
        server.user_message(user_id=user_id, agent_id=agent_state.id, message="/memory")
        raise Exception("user_message call should have failed")
    except ValueError as e:
        # Error is expected
        print(e)
    except:
        raise

    print(server.run_command(user_id=user_id, agent_id=agent_state.id, command="/memory"))

    # add data into archival memory
    agent = server._load_agent(user_id=user_id, agent_id=agent_state.id)
    archival_memories = ["alpha", "Cinderella wore a blue dress", "Dog eat dog", "ZZZ", "Shishir loves indian food"]
    embed_model = embedding_model(embedding_config)
    for text in archival_memories:
        embedding = embed_model.get_text_embedding(text)
        agent.persistence_manager.archival_memory.storage.insert(
            Passage(user_id=user_id, agent_id=agent_state.id, text=text, embedding=embedding)
        )

    # add data into recall memory
    server.user_message(user_id=user_id, agent_id=agent_state.id, message="Hello?")
    server.user_message(user_id=user_id, agent_id=agent_state.id, message="Hello?")
    server.user_message(user_id=user_id, agent_id=agent_state.id, message="Hello?")
    server.user_message(user_id=user_id, agent_id=agent_state.id, message="Hello?")
    server.user_message(user_id=user_id, agent_id=agent_state.id, message="Hello?")

    # test recall memory cursor pagination
    cursor1, messages_1 = server.get_agent_recall_cursor(user_id=user_id, agent_id=agent_state.id, reverse=True, limit=2)
    cursor2, messages_2 = server.get_agent_recall_cursor(user_id=user_id, agent_id=agent_state.id, reverse=True, after=cursor1, limit=1000)
    cursor3, messages_3 = server.get_agent_recall_cursor(user_id=user_id, agent_id=agent_state.id, reverse=True, limit=1000)
    ids3 = [m["id"] for m in messages_3]
    ids2 = [m["id"] for m in messages_2]
    timestamps = [m["created_at"] for m in messages_3]
    print("timestamps", timestamps)
    assert messages_3[-1]["created_at"] < messages_3[0]["created_at"]
    assert len(messages_3) == len(messages_1) + len(messages_2)
    cursor4, messages_4 = server.get_agent_recall_cursor(user_id=user_id, agent_id=agent_state.id, reverse=True, before=cursor1)
    assert len(messages_4) == 1

    # test archival memory cursor pagination
    cursor1, passages_1 = server.get_agent_archival_cursor(
        user_id=user_id, agent_id=agent_state.id, reverse=False, limit=2, order_by="text"
    )
    cursor2, passages_2 = server.get_agent_archival_cursor(
        user_id=user_id, agent_id=agent_state.id, reverse=False, after=cursor1, order_by="text"
    )
    cursor3, passages_3 = server.get_agent_archival_cursor(
        user_id=user_id, agent_id=agent_state.id, reverse=False, before=cursor2, limit=1000, order_by="text"
    )
    print("p1", [p["text"] for p in passages_1])
    print("p2", [p["text"] for p in passages_2])
    print("p3", [p["text"] for p in passages_3])
    assert passages_1[0]["text"] == "alpha"
    assert len(passages_2) == 3
    assert len(passages_3) == 4

    # test recall memory
    messages_1 = server.get_agent_messages(user_id=user_id, agent_id=agent_state.id, start=0, count=1)
    assert len(messages_1) == 1
    messages_2 = server.get_agent_messages(user_id=user_id, agent_id=agent_state.id, start=1, count=1000)
    messages_3 = server.get_agent_messages(user_id=user_id, agent_id=agent_state.id, start=1, count=5)
    # not sure exactly how many messages there should be
    assert len(messages_2) > len(messages_3)
    # test safe empty return
    messages_none = server.get_agent_messages(user_id=user_id, agent_id=agent_state.id, start=1000, count=1000)
    assert len(messages_none) == 0

    # test archival memory
    passage_1 = server.get_agent_archival(user_id=user_id, agent_id=agent_state.id, start=0, count=1)
    assert len(passage_1) == 1
    passage_2 = server.get_agent_archival(user_id=user_id, agent_id=agent_state.id, start=1, count=1000)
    assert len(passage_2) == 4
    # test safe empty return
    passage_none = server.get_agent_archival(user_id=user_id, agent_id=agent_state.id, start=1000, count=1000)
    assert len(passage_none) == 0


if __name__ == "__main__":
    test_server()