import json
import logging
import os
import sys
from datetime import datetime, timezone
from pathlib import Path

from dotenv import load_dotenv
from rich.pretty import pprint

from graphiti_core import Graphiti
from graphiti_core.edges import EntityEdge
from graphiti_core.llm_client.anthropic_client import AnthropicClient
from graphiti_core.nodes import EpisodeType
from graphiti_core.utils.bulk_utils import RawEpisode
from graphiti_core.utils.maintenance.graph_data_operations import clear_data

load_dotenv()

neo4j_uri = os.environ.get('NEO4J_URI', 'bolt://localhost:7687')
neo4j_user = os.environ.get('NEO4J_USER', 'neo4j')
neo4j_password = os.environ.get('NEO4J_PASSWORD', 'password')

def setup_logging():
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)
    formatter = logging.Formatter('%(name)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    return logger


logger = setup_logging()

shoe_conversation_1 = [
    "SalesBot (2024-07-30T00:00:00Z): Hi, I'm ManyBirds Assistant! How can I help you today?",
    "John (2024-07-30T00:01:00Z): Hi, I'm looking for a new pair of shoes.",
    'SalesBot (2024-07-30T00:02:00Z): Of course! What kind of material are you looking for?',
    "John (2024-07-30T00:03:00Z): I'm allergic to wool. Also, I'm a size 10 if that helps?",
    "SalesBot (2024-07-30T00:04:00Z): We have just what you are looking for, how do you like our Men's Couriers. They have a retro silhouette look and from cotton. How about them in Basin Blue?",
    "John (2024-07-30T00:05:00Z): Blue is great! Love the look. I'll take them.",
]

shoe_conversation_2 = [
    'SalesBot (2024-08-20T00:00:00Z): Hi John, how can I assist you today?',
    "John (2024-08-20T00:01:00Z): Hi, I need to return the Men's Couriers I bought recently. They're too tight for my wide feet. Hahaha.",
    "SalesBot (2024-08-20T00:02:00Z): I'm sorry to hear that. We can process the return for you.",
]

async def add_messages(client: Graphiti, messages: list[str], prefix: str = 'Message'):
    for i, message in enumerate(messages):
        await client.add_episode(
            name=f'{prefix}-{i}',
            episode_body=message,
            source=EpisodeType.message,
            reference_time=datetime.now(timezone.utc),
            source_description='Shoe conversation',
        )

async def ingest_products_data(client: Graphiti):
    script_dir = Path.cwd().parent
    json_file_path = script_dir / 'data' / 'manybirds_products.json'

    with open(json_file_path) as file:
        products = json.load(file)['products']

    episodes: list[RawEpisode] = [
        RawEpisode(
            name=product.get('title', f'Product {i}'),
            content=str({k: v for k, v in product.items() if k != 'images'}),
            source_description='ManyBirds products',
            source=EpisodeType.json,
            reference_time=datetime.now(timezone.utc),
        )
        for i, product in enumerate(products)
    ]

    await client.add_episode_bulk(episodes)

def pretty_print(entity: EntityEdge | list[EntityEdge]):
    if isinstance(entity, EntityEdge):
        data = {k: v for k, v in entity.model_dump().items() if k != 'fact_embedding'}
    elif isinstance(entity, list):
        data = [{k: v for k, v in e.model_dump().items() if k != 'fact_embedding'} for e in entity]
    else:
        pprint(entity)
        return
    pprint(data)

llm_client = AnthropicClient(cache=False)

client = Graphiti(
    neo4j_uri,
    neo4j_user,
    neo4j_password,
    llm_client=llm_client,
)

await clear_data(client.driver)
await client.build_indices_and_constraints()
await ingest_products_data(client)

await client.add_episode(
    name='Inventory management 0',
    episode_body=('All Tinybirds Wool Runners styles are out of stock until December 25th 2024'),
    source=EpisodeType.text,
    reference_time=datetime.now(timezone.utc),
    source_description='Inventory Management Bot',
)

r = await client.search('Which products are out of stock?')

pretty_print(r[0])

await add_messages(client, shoe_conversation_1, prefix='conversation-1')

r = await client.search("What is John's shoe size?", num_results=2)

pretty_print(r)

from graphiti_core.search.search_config_recipes import NODE_HYBRID_SEARCH_RRF

nl = await client._search('John', NODE_HYBRID_SEARCH_RRF)

pretty_print(nl[0])

john_uuid = nl[0].uuid

r = await client.search('Can John wear ManyBirds Wool Runners?', num_results=3)

print('-' * 100)
print('Standard Reciprocal Rank Fusion Reranking')
print('-' * 100)
for record in r:
    print(record.fact)

r = await client.search(
    'Can John wear ManyBirds Wool Runners?', center_node_uuid=john_uuid, num_results=3
)

print('-' * 100)
print("Node Distance Reranking from 'John' node")
print('-' * 100)
for record in r:
    print(record.fact)

await add_messages(client, shoe_conversation_2, prefix='conversation-2')

r = await client.search('What shoes has John purchased?', center_node_uuid=john_uuid, num_results=3)

pretty_print(r)

r = await client.search('What shoes has John purchased?', center_node_uuid=john_uuid, num_results=5)

pretty_print(r)

r = await client.search('Who is John?', num_results=5)

pretty_print(r)

r = await client.search(
    'What did John do about his discomfort with the Mens Couriers shoes', num_results=5
)

pretty_print(r)