import sqlite3
import pandas as pd
from mcp_memory_search.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 mcp_memory_search.tool_schema import MEMORY_SEARCH_TOOLS
from mcp_memory_search 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_SEARCH_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 == "basic_memory_search":
            return await basic_memory_search(arguments or {})
        elif name == "advanced_memory_search":
            return await advanced_memory_search(arguments or {})
        elif name == "get_recent_conversations":
            return await get_recent_conversations(arguments or {})
        elif name == "get_session_history":
            return await get_session_history(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:
def parse_time_scope(time_scope: str) -> Optional[datetime]:
    """Convert time scope to reference datetime."""
    now = datetime.now(timezone.utc)
    
    if time_scope == "recent":
        return now - timedelta(days=30)
    elif time_scope == "very_recent":
        return now - timedelta(days=7)
    else:  # all_time
        return None


def get_advanced_search_config(ranking_strategy: str, variety_level: float = 0.5):
    """Map ranking strategy to search configuration with variety control."""
    try:
        from graphiti_core.search.search_config_recipes import (
            COMBINED_HYBRID_SEARCH_CROSS_ENCODER,
            COMBINED_HYBRID_SEARCH_RRF,
            COMBINED_HYBRID_SEARCH_MMR,
            EDGE_HYBRID_SEARCH_NODE_DISTANCE,
            EDGE_HYBRID_SEARCH_EPISODE_MENTIONS,
            NODE_HYBRID_SEARCH_NODE_DISTANCE,
        )
        
        if ranking_strategy == "most_accurate":
            return COMBINED_HYBRID_SEARCH_CROSS_ENCODER
        elif ranking_strategy == "proximity_focused":
            return EDGE_HYBRID_SEARCH_NODE_DISTANCE
        elif ranking_strategy == "recent_mentions":
            return EDGE_HYBRID_SEARCH_EPISODE_MENTIONS
        elif ranking_strategy == "diverse_perspectives":
            return COMBINED_HYBRID_SEARCH_MMR
        else:  # balanced_ranking (default)
            return COMBINED_HYBRID_SEARCH_RRF
    except ImportError:
        return None
    


def get_search_filters(time_constraints: dict | None = None):
    """Create search filters from user parameters."""
    try:
        from graphiti_core.search.search_filters import SearchFilters, DateFilter, ComparisonOperator
        
        filters = SearchFilters()
        
        if time_constraints:
            date_filters = []
            
            if time_constraints.get('start_date'):
                start_date = datetime.fromisoformat(time_constraints['start_date'].replace('Z', '+00:00'))
                start_filter = DateFilter(
                    date=start_date,
                    comparison_operator=ComparisonOperator.greater_than_equal
                )
                date_filters.append(start_filter)
            
            if time_constraints.get('end_date'):
                end_date = datetime.fromisoformat(time_constraints['end_date'].replace('Z', '+00:00'))
                end_filter = DateFilter(
                    date=end_date,
                    comparison_operator=ComparisonOperator.less_than_equal
                )
                date_filters.append(end_filter)
            
            if date_filters:
                filters.created_at = [date_filters]  # All conditions must be met (AND)
                filters.valid_at = [date_filters]
            
        return filters
    except ImportError:
        return None
    
def get_session(session_id: str) -> SessionDetails | ErrorResponse:
    conn = sqlite3.connect(PATH_TO_DB)
    try:
        # Identify all chat-like tables
        table_names = [row[0] for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table';").fetchall()]
        chat_tables = [t for t in table_names if "chat" in t.lower()]

        if not chat_tables:
            raise ValueError("❌ No chat-related tables found in the database.")

        all_rows = []

        # Collect data across all chat tables
        for table in chat_tables:
            df = pd.read_sql_query(
                f"SELECT question, answer, createdAt FROM {table} WHERE session_id = ?",
                conn,
                params=(session_id,)
            )

            if not df.empty:
                df['createdAt'] = pd.to_datetime(df['createdAt'], errors='coerce')
                all_rows.append(df)

        if not all_rows:
           return SessionDetails(session_id=session_id, session_created_at= "", session_updated_at= "", messages=[])

        # Combine and sort all rows by createdAt
        df_all = pd.concat(all_rows)
        df_all = df_all.dropna(subset=['createdAt']).sort_values('createdAt')

        message_list = json.loads(df_all[['question', 'answer']].to_json(orient='records'))
        created_at = df_all['createdAt'].iloc[0].to_pydatetime().isoformat()
        updated_at = df_all['createdAt'].iloc[-1].to_pydatetime().isoformat()


        return SessionDetails(session_id=session_id, session_created_at=created_at, session_updated_at=updated_at, messages=message_list)
    
    except Exception as e:
        logger.error(f"Error getting session history: {e}")
        return ErrorResponse(error=str(e))


def format_search_results(results) -> str:
    """
    Format search results to a JSON string, ensuring all embeddings are excluded.
    """
    if not results:
        return json.dumps({"message": "No results found"}, indent=2)

    output = {}

    # Process nodes, edges, episodes, and communities if they exist
    if hasattr(results, 'nodes') and results.nodes:
        output['nodes'] = []
        for node in results.nodes:
            node_dict = node.model_dump(mode='json', exclude={'node_embedding', 'name_embedding'})
            if 'attributes' in node_dict and isinstance(node_dict['attributes'], dict):
                node_dict['attributes'].pop('node_embedding', None)
                node_dict['attributes'].pop('name_embedding', None)
            output['nodes'].append(node_dict)

    if hasattr(results, 'edges') and results.edges:
        output['edges'] = []
        for edge in results.edges:
            edge_dict = edge.model_dump(mode='json', exclude={'fact_embedding'})
            if 'attributes' in edge_dict and isinstance(edge_dict['attributes'], dict):
                edge_dict['attributes'].pop('fact_embedding', None)
            output['edges'].append(edge_dict)

    if hasattr(results, 'episodes') and results.episodes:
        output['conversations'] = [
            episode.model_dump(mode='json', exclude={'episode_embedding'})
            for episode in results.episodes
        ]
    
    if hasattr(results, 'communities') and results.communities:
        output['communities'] = [
            community.model_dump(mode='json', exclude={'community_embedding'})
            for community in results.communities
        ]

    # Clean up empty lists from the final output
    final_output = {k: v for k, v in output.items() if v}

    return json.dumps(final_output, indent=2)
    

#actual tools:

async def basic_memory_search(arguments: dict[str, Any]) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Perform a smart search across memory, allowing for filtering and focus.
    """
    try:
        query = arguments.get("query")
        if not query:
            raise ValueError("Query is required for basic memory search.")
        group_ids = arguments.get("group_ids", [])
        search_focus = arguments.get("search_focus", "comprehensive")
        time_scope = arguments.get("time_scope", "all_time")
        max_results = arguments.get("max_results", 15)
        center_around = arguments.get("center_around")

        graphiti_client = await initialize_graphiti()

        config = COMBINED_HYBRID_SEARCH_RRF.model_copy(deep=True)
        config.limit = max_results
        
        # Handle time scope filtering for basic search
        search_filter = None
        if time_scope in ["recent", "very_recent"]:
            try:
                from graphiti_core.search.search_filters import SearchFilters, DateFilter, ComparisonOperator
                
                reference_time = parse_time_scope(time_scope)
                if reference_time:
                    search_filter = SearchFilters()
                    date_filter = DateFilter(
                        date=reference_time,
                        comparison_operator=ComparisonOperator.greater_than_equal
                    )
                    search_filter.valid_at = [[date_filter]]
                    search_filter.created_at = [[date_filter]]
            except ImportError:
                pass
      

        results = await graphiti_client.search_(
            query=query,
            config=config,
            group_ids=group_ids,
            center_node_uuid=center_around,
            search_filter=search_filter
        )

        # Filter results based on focus
        if search_focus == "facts_only":
            results.nodes = []
            results.episodes = []
        elif search_focus == "entities_only":
            results.edges = []
            results.episodes = []
        elif search_focus == "conversations_only":
            results.edges = []
            results.nodes = []

        return [types.TextContent(text=format_search_results(results), type="text")]
    except Exception as e:
        logger.error(f"Error in smart memory search: {e}")
        return [types.TextContent(text=f"Error: {str(e)}", type="text")] 
    


async def advanced_memory_search(arguments: dict[str, Any]) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:

    query = arguments["query"]
    if not query:
        raise ValueError("Query is required for priority-based memory search.")
    group_ids = arguments.get("group_ids", [])
    ranking_strategy = arguments.get("ranking_strategy", "most_accurate")
    search_focus = arguments.get("search_focus", "comprehensive")
    reference_point = arguments.get("reference_point")
    relationship_focus = arguments.get("relationship_focus", [])
    time_constraints = arguments.get("time_constraints", {})
    max_results = arguments.get("max_results", 25)

    try :
        graphiti_client = await initialize_graphiti()

        config = get_advanced_search_config(ranking_strategy)
        if not config:
            config = COMBINED_HYBRID_SEARCH_RRF.model_copy(deep=True)
        config.limit = max_results

        search_filter = get_search_filters(time_constraints)

        results = await graphiti_client.search_(
            query=query,
            config=config,
            group_ids=group_ids,
            center_node_uuid=reference_point,
            search_filter=search_filter
        )

        if not results:
            return [types.TextContent(text="No results found", type="text")]

        # Filter results based on focus
        if search_focus == "facts_only":
            results.nodes = []
            results.episodes = []
        elif search_focus == "entities_only":
            results.edges = []
            results.episodes = []
        elif search_focus == "conversations_only":
            results.edges = []
            results.nodes = []

        return [types.TextContent(text=format_search_results(results), type="text")]
    except Exception as e:
        logger.error(f"Error in advanced memory search: {e}")
        return [types.TextContent(text=f"Error: {str(e)}", type="text")]
    


async def get_recent_conversations(arguments: dict[str, Any]) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    last_n = arguments.get("last_n", 10)
    group_ids = arguments.get("group_ids", [])
    
    try:
        graphiti_client = await initialize_graphiti()
        episodes = await graphiti_client.retrieve_episodes(group_ids=group_ids, last_n=last_n, reference_time=datetime.now(timezone.utc))
        if not episodes:
            return [types.TextContent(text="No episodes found", type="text")]
        
        formatted_episodes = [
            episode.model_dump(mode='json')
            for episode in episodes
        ]
        return [types.TextContent(text=json.dumps(formatted_episodes, indent=2), type="text")]
    except Exception as e:
        logger.error(f"Error getting episodes: {e}")
        return [types.TextContent(text=f"Error: {str(e)}", type="text")]
    

async def get_session_history(arguments: dict[str, Any]) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    session_id = arguments.get("session_id")
    if not session_id:
        raise ValueError("Session ID is required for session history search.")
    
    try:
        session_history = get_session(session_id)
        return [types.TextContent(text=json.dumps(session_history, indent=2), type="text")]
    except Exception as e:
        logger.error(f"Error getting session history: {e}")
        return [types.TextContent(text=f"Error: {str(e)}", type="text")]
    