import re
from _arango import arango
from _llm import LLM
from pprint import pprint
from pprint import pprint
from langchain_text_splitters import CharacterTextSplitter
import multiprocessing
from print_color import *


def describe_relation(person1, person2, relation, text):
    llm = LLM(chat=False, system_prompt="Du ska hitta relationer i en text. Svara alltid enligt angiven form och alltid på svenska.")
    prompt = f'''
    I texten nedan beskrivs att {person1} och {person2} har relationen "{relation}". Läs texten och sammanfatta kortfattat vad som beskrivs om relationen mellan {person1} och {person2}:\n\n"""{text}"""\n
    Svara ENBART med information om relationen, inga hälsningsfraser eller liknande.
    Relationen ska vara kortfattad och stämma med texten. Om det inte går att beskriva relationen så svara med "None".
    Vad står det om relationen "{relation}" mellan {person1} och {person2}?
    '''
    response = llm.generate(prompt)
    print_rainbow(relation, response)
    return response


def find_relations(interrogation):
    """
    Finds the relations in an interrogation.

    Args:
        interrogation (dict): The interrogation.

    Returns:
        None
    """
    text = interrogation["text"]
    text_splitter = CharacterTextSplitter(
        separator="\n\n",
        chunk_size=6000,
        chunk_overlap=0,
        length_function=len,
        is_separator_regex=False,
    )
    chunks = text_splitter.split_text(text)
    all_relations = []
    llm = LLM(chat=False, system_prompt="Du ska hitta relationer i en text. Svara alltid enligt angiven form och alltid på svenska.")

    for chunk in chunks:
        prompt = f"""Nedan är en bit av ett polisförhör med {interrogation['person']}. Jag vill att du hittar alla relationer mellan identifierbara personer som beskrivs i själva förhöret:\n\n{chunk}\n\n
        Svara på formen "person1;person2;relation\n". Var noga med hur semikolon används för att skilja personerna och relationen, och ny rad efter varje relation (informationen ska sedan användas för en CSV fil). Svara på svenska.
        Nedan är ett påhittat exempel för att du ska förstå hur du kan svara:
        <exempel>
        person1;person2;gick på grundskolan tillammans, spelade fotboll
        person2;person3;gifta sedan 2022
        </exempel>
        Beskrivningen av relationen ska vara kortfattad och stämma med texten.
        Om det inte finns någon relation, svara "None". 
        Svara ENBART med relationerna, INGENTING annat som en hälsning eller förklaring.
        """
        response = llm.generate(prompt)
        print_blue(response)
        relations = response.split("\n")
        for relation in relations:
            if relation == "None" or ';' not in relation:
                continue
            try:
                person1, person2, relation = relation.split(";", 2)
            except ValueError as e:
                print_red(f"Error: {e}")
                print_red(f"Relation: {relation}")
                continue
            
            description = describe_relation(person1, person2, relation, chunk)
            

            for r in all_relations:
                p1 = r["from"]
                p2 = r["to"]
                if p1 == person1 and p2 == person2:
                    r["relations"].append({'relation': relation, 'description': description, 'chunk_number': chunks.index(chunk)})
                    break
            else:
                all_relations.append(
                    {   
                        "_key": arango.fix_key_name(f"{person1}-{person2}_{interrogation['_key']}"),
                        "from": person1,
                        "to": person2,
                        "relations": [{'relation': relation, 'description': description, 'chunk_number': chunks.index(chunk)}],
                        'interrogation': interrogation['_key'],
                        "chunks": chunks,
                    }
                )
                

    return all_relations

if __name__ == "__main__":
    db = arango.db

    q = 'for doc in interrogations return doc'
    interrogations = list(db.aql.execute(q))
    for interrogation in interrogations:
        relations = find_relations(interrogation)
        db.collection('relations').insert_many(relations)