from datetime import datetime, timedelta

from fastapi import FastAPI

from mem0x.app.instance import memoryInstance
from mem0x.app.type import (BootstrapRequest, LongTermAddRequest,
                            LongTermRemoveAllRequest, LongTermRemoveRequest,
                            LongTermSearchRequest, LongTermUpdateRequest,
                            Response, ShortTermAddRequest,
                            ShortTermDeleteRequest, ShortTermGetRequest,
                            ShortTermGetRunIdsRequest)
from mem0x.bootstrap.config import AddConfig, ChunkerConfig
from mem0x.bootstrap.data_formatter import DataFormatter
from mem0x.bootstrap.models.data_type import DataType
from mem0x.configs.prompts import (FACT_RETRIEVAL_PROMPT,
                                   FACT_RETRIEVAL_PROMPT_KNOWLEDGE)

app = FastAPI()


@app.post("/longterm_search")
async def longterm_search(search_request: LongTermSearchRequest):
    try:
        filters = {"table_name": search_request.table_name}
        related_memories = memoryInstance.m.search(
            search_request.query,
            search_request.user_id,
            search_request.agent_id,
            search_request.run_id,
            search_request.limit,
            filters,
        )
        return Response(message="SUCCESS", status=True, data=related_memories)
    except Exception as e:
        return Response(message=f"{e}", status=False, data=None)


@app.post("/longterm_add")
async def longterm_add(add_request: LongTermAddRequest):
    try:
        print(add_request.messages)
        metadata = {"table_name": add_request.table_name}
        filters = {"table_name": add_request.table_name}
        result = memoryInstance.m.add(
            add_request.messages,
            add_request.user_id,
            add_request.agent_id,
            add_request.run_id,
            metadata,
            filters,
            FACT_RETRIEVAL_PROMPT,
        )
        return Response(message="SUCCESS", status=True, data=result)
    except Exception as e:
        return Response(message=f"{e}", status=False, data=None)


@app.post("/longterm_remove")
async def longterm_remove(remove_request: LongTermRemoveRequest):
    try:
        memoryInstance.m.delete(remove_request.memory_id)
        return Response(message="SUCCESS", status=True, data=None)
    except Exception as e:
        return Response(message=f"{e}", status=False, data=None)


@app.post("/longterm_remove_all")
async def longterm_remove_all(remove_all_request: LongTermRemoveAllRequest):
    try:
        memoryInstance.m.delete_all(remove_all_request.user_id, remove_all_request.agent_id, remove_all_request.run_id)
        return Response(message="SUCCESS", status=True, data=None)
    except Exception as e:
        return Response(message=f"{e}", status=False, data=None)


@app.post("/longterm_update")
async def longterm_update(update_request: LongTermUpdateRequest):
    try:
        memoryInstance.m.update(update_request.memory_id, update_request.data)
        return Response(message="SUCCESS", status=True, data=None)
    except Exception as e:
        return Response(message=f"{e}", status=False, data=None)


@app.post("/bootstrap_add")
async def bootstrap_add(bootstrap_request: BootstrapRequest):
    try:
        loader, chunker = None, None
        config = AddConfig()
        config.chunker = ChunkerConfig(
            chunk_size=bootstrap_request.chunk_size,
            chunk_overlap=bootstrap_request.chunk_overlap,
            min_chunk_size=bootstrap_request.min_chunk_size,
        )
        data_type = DataType(bootstrap_request.data_type)
        data_formatter = DataFormatter(data_type, config, loader, chunker)

        src = bootstrap_request.src
        embeddings_data = data_formatter.chunker.create_chunks(data_formatter.loader, src)

        messages = []
        for document in embeddings_data["documents"]:
            messages.append({"role": "", "content": document})

        metadata = {"table_name": bootstrap_request.table_name}
        filters = {"table_name": bootstrap_request.table_name}
        result = memoryInstance.m.add(
            messages,
            bootstrap_request.user_id,
            bootstrap_request.agent_id,
            bootstrap_request.run_id,
            metadata,
            filters,
            FACT_RETRIEVAL_PROMPT_KNOWLEDGE,
        )
        return Response(message="SUCCESS", status=True, data=result)
    except Exception as e:
        return Response(message=f"{e}", status=False, data=None)


@app.post("/shortterm_add")
async def shortterm_add(add_request: ShortTermAddRequest):
    try:
        memoryInstance.sm.add(
            add_request.content, add_request.user_id, add_request.agent_id, add_request.run_id, add_request.table_name
        )
        return Response(message="SUCCESS", status=True, data=None)
    except Exception as e:
        return Response(message=f"{e}", status=False, data=None)


@app.post("/shortterm_get")
async def shortterm_get(get_request: ShortTermGetRequest):
    try:
        timestamp = get_request.timestamp
        if timestamp is None:
            comparison_time = datetime.now() + timedelta(hours=24)
            timestamp = comparison_time.isoformat()

        result = memoryInstance.sm.get(get_request.table_name, get_request.run_id, timestamp)
        return Response(message="SUCCESS", status=True, data=result)
    except Exception as e:
        return Response(message=f"{e}", status=False, data=None)


@app.post("/shortterm_get_runids")
async def shortterm_get_runids(get_runids_request: ShortTermGetRunIdsRequest):
    try:
        timestamp = get_runids_request.timestamp
        if timestamp is None:
            comparison_time = datetime.now() + timedelta(hours=24)
            timestamp = comparison_time.isoformat()

        result = memoryInstance.sm.get_runids(timestamp)
        return Response(message="SUCCESS", status=True, data=result)
    except Exception as e:
        return Response(message=f"{e}", status=False, data=None)


@app.post("/shortterm_delete")
async def shortterm_delete(delete_request: ShortTermDeleteRequest):
    try:
        memoryInstance.sm.delete(delete_request.id)
        return Response(message="SUCCESS", status=True, data=None)
    except Exception as e:
        return Response(message=f"{e}", status=False, data=None)
