from _chroma import ChromaDB
from _arango import arango
from _llm import LLM
from pprint import pprint
from print_color import *
import multiprocessing

def add_persons_to_chroma():
    db = arango.db
    q = "for doc in persons filter doc.other != true return doc"
    persons = list(db.aql.execute(q))

    # Lists to store the documents, metadatas and ids
    documents = []
    metadatas = []
    ids = []

    for person in persons:
        if 'info' in person:
            info = '\n'.join(person['info'])
            documents.append(person['name'])
            #documents.append(f"{person['name']}\n{info}")
            metadata = {'name': person['name'], '_key': person['_key']}
            metadatas.append(metadata)
            ids.append(person["_key"])

    collection = chroma.client.get_collection('mala_persons')
    collection.add(documents=documents, metadatas=metadatas, ids=ids)


def find_person(person):
    """
    Finds a person in the database based on the given person information.

    Args:
        person (dict): A dictionary containing information about the person.

    Returns:
        tuple: A tuple containing the following elements:
            - answer (str): The generated answer from the language model.
            - person_in_arango (dict): The information of the person retrieved from ArangoDB.
            - interrogation (str): The ID of the interrogation where the person was mentioned.
    """
    db = arango.db
    llm = LLM()

    found = False
    other_person = person['name']

    chroma = ChromaDB()
    col = chroma.client.get_or_create_collection('mala_persons')
    hits = col.query(query_texts=[other_person], n_results=1)

    hit = hits['documents'][0][0]
    _key = hits['metadatas'][0][0]['_key']
    distance = hits['distances'][0][0]
    
    #* Filter out hits with distance > 1
    if distance > 1:
        return None
    person_in_arango = db.collection('persons').get(_key)
    info = '\n'.join(person_in_arango['info'])

    print('Other:', other_person, 'Förslag:', hit, 'Distance:', distance)
    interrogations = person['mentioned_in_interrogation']

    output = []
    for interrogation in interrogations:
        text = db.collection('interrogations').get(interrogation)['text']
        prompt = f'''I texten nedan omnämns en "{other_person}" och jag försöker förstå om det kan vara exempelvis ett felstavat namn eller smeknamn för en annan person.\n
        TEXT:
        """{text}"""\n

        På andra ställen i polisens förundersökning finns en person som heter "{hit}", och som beskrivs så här:
        """{info}"""\n
        Verkar det troligt att personen som kallas {other_person} är samma person som {hit}? Svara bara JA eller NEJ, samt en kort förklaring till varför.
        '''
        #  Om istället förnamnet eller efternamnet är helt olika så är det förmodligen inte samma person.Om det bara är ett namn (inget efternamn) kan det också handla om ett smeknamn eller en beskrivning.
        answer = llm.generate(prompt)

        output.append((answer, person_in_arango, text))
    return output

def verify(answer, person, person_in_arango, text):
    found_input = input('Enter om det stämmer? >> ')
    if found_input == '':
        found = True
        person['mentioned_in_interrogation'].remove(text)
        db.collection('persons').update(person)
        person_in_arango['info'].append(answer)
        person_in_arango['mentioned_in_interrogation'].append(text)
        db.collection('persons').update(person_in_arango)
    if found:
        db.collection('other_persons').insert(person, overwrite=True)
        db.collection('persons').delete(person, check_rev=False)
        
if __name__ == '__main__':
    db = arango.db
    persons = list(db.collection('persons').all())
    
    q = 'for doc in persons filter doc.other == true return doc'
    other_persons = [i for i in db.aql.execute(q)]

    for person in other_persons:
        print(find_person(person))
        exit()
    # with multiprocessing.Pool() as pool:
    #     pool.map(find_person, other_persons)