import json
import logging
import time
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()

logger = logging.getLogger(__name__)


@app.post("/longterm_search")
async def longterm_search(search_request: LongTermSearchRequest):
    try:
        logging.info(f"/longterm_search?search_request={search_request.model_dump_json()}")
        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:
        logging.info(f"/longterm_add?add_request={add_request.model_dump_json()}")
        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:
        logging.info(f"/longterm_remove?remove_request={remove_request.model_dump_json()}")
        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:
        logging.info(f"/longterm_update?update_request={update_request.model_dump_json()}")
        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:
        logging.info(f"/bootstrap_request?bootstrap_request={bootstrap_request.model_dump_json()}")
        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(document)

        logging.info(f"documents size={len(messages)}")
        metadata = {"table_name": bootstrap_request.table_name}
        filters = {"table_name": bootstrap_request.table_name}
        result = []
        for document in messages:
            r = memoryInstance.m.add(
                document,
                bootstrap_request.user_id,
                bootstrap_request.agent_id,
                bootstrap_request.run_id,
                metadata,
                filters,
                FACT_RETRIEVAL_PROMPT_KNOWLEDGE,
            )
            result.append(r)
            logging.info(f"r={r}")
            time.sleep(3)
        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:
        logging.info(f"/shortterm_add?add_request={add_request.model_dump_json()}")
        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:
        logging.info(f"/shortterm_get?get_request={get_request.model_dump_json()}")
        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:
        logging.info(f"/shortterm_get_runids?get_runids_request={get_runids_request.model_dump_json()}")
        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:
        logging.info(f"/shortterm_delete?delete_request={delete_request.model_dump_json()}")
        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)
