import sqlite3
import pandas as pd
from mcp_memory_builder.constants import *
import json
from datetime import datetime, timezone
from graphiti_core.graphiti import Graphiti
from graphiti_core.search.search_config_recipes import (
    COMBINED_HYBRID_SEARCH_RRF
)
from typing import Optional, Any, TypedDict, Sequence, Union
from datetime import timedelta
import mcp.types as types
from graphiti_core.llm_client import LLMConfig, OpenAIClient
from graphiti_core.nodes import EpisodeType



from mcp_memory_builder.tool_schema import MEMORY_BUILDER_TOOLS
from mcp_memory_builder import mcp,logger




# Configure to use GPT-4o (more powerful model)
llm_config = LLMConfig(
    model=MODEL_NAME,  # Options: "gpt-4o", "gpt-4o-mini", "gpt-4-turbo"
    temperature=0.0,  # Lower temperature for more consistent extraction
    max_tokens=4000
)
llm_client = OpenAIClient(llm_config)


def register_tools():
    @mcp.list_tools()
    async def handle_list_tools() -> list[types.Tool]:
        return MEMORY_BUILDER_TOOLS
    @mcp.call_tool()
    async def handle_call_tool(
        name: str, arguments: dict[str, Any] | None
    ) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
        if name == "add_memory":
            return await add_memory(arguments or {})
        else:
            raise ValueError(f"Unknown tool: {name}")


class SessionDetails(TypedDict):
    session_id: str
    session_created_at: str
    session_updated_at: str
    messages: list[dict[str, Any]]

class ErrorResponse(TypedDict):
    error: str
async def initialize_graphiti():

    try:
        graphiti_client = Graphiti(
                uri = NEO4J_URI,
                user= NEO4J_USER,
                password= NEO4J_PASSWORD,
                llm_client=llm_client
                )
        await graphiti_client.build_indices_and_constraints()
        logger.info("Graphiti client initialized")

        return graphiti_client
    except Exception as e:
        logger.error(f"Error initializing Graphiti client: {e}")
        raise


#utility functions:
async def add_memory(arguments: dict[str, Any]) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Add a memory to the graphiti client.
    """
    name = arguments.get("name", None)
    content = arguments.get("content", None)
    source_description = arguments.get("source_description", None)
    content_type = arguments.get("content_type", None)
    reference_time = arguments.get("reference_time",None)
    group_id = arguments.get("group_id",None)
    session_id = arguments.get("session_id",None)
    update_communities = arguments.get("update_communities", False)
    metadata = arguments.get("metadata", {})
    previous_episode_uuids = arguments.get("previous_episode_uuids", None)

    if not name or not content or not source_description or not content_type:
        return [types.TextContent(text="Missing required fields", type="text")]
    
    if content_type not in ["text", "json", "messages"]:
        return [types.TextContent(text="Invalid content type", type="text")]
    

    if content_type == "messages":
        source = EpisodeType.message
    elif content_type == "json":
        source = EpisodeType.json
    elif content_type == "text":
        source = EpisodeType.text

    if not source:
        source = EpisodeType.message

    if not reference_time:
        reference_time = datetime.now(timezone.utc)
    
    if not group_id:
       return [types.TextContent(text="Group ID is required", type="text")]
    
    if not session_id:
        return [types.TextContent(text="Session ID is required", type="text")]
    
    try:
        graphiti_client = await initialize_graphiti()
    
        result = await graphiti_client.add_episode(
            name=name,
            episode_body=content,
            source_description=source_description,
            reference_time=reference_time,
            source=source,
            group_id=group_id,
            session_id=session_id,
            update_communities=update_communities,
            metadata=metadata,
            previous_episode_uuids=previous_episode_uuids
        )

        return [types.TextContent(text=f"Memory added successfully to episode: {result.episode.uuid}", type="text")]

    except Exception as e:
        return [types.TextContent(text=f"Error adding memory: {e}", type="text")]
    
    
    
    
