import re
from _arango import arango
from _llm import LLM

from pprint import pprint
from langchain_text_splitters import CharacterTextSplitter

# Create an instance of the SentenceSplitter
text_splitter = CharacterTextSplitter(
    separator="\n\n",
    chunk_size=6000,
    chunk_overlap=0,
    length_function=len,
    is_separator_regex=False,
)

llm = LLM(chat=False)

interrogations = [i for i in arango.db.collection("interrogations").all()]
interrogations = sorted(interrogations, key=lambda x: x["date"])


# Sort out interrogations already processed
all_relations = [i for i in arango.db.collection("all_relations").all()]
for relation in all_relations:
    if "interrogation" in relation["context"]:
        for interrogation in interrogations:
            if relation["_in"] == interrogation["_id"]:
                interrogations.remove(interrogation)


# Go through all interrogations
for interrogation in interrogations:

    # Get the persons (now updated from the last one, so it should be all persons in the database)
    persons_docs = []
    for collection in ["persons", "other_persons"]:
        for i in arango.db.collection(collection).all():
            persons_docs.append(i)
    persons = [i["name"].strip() for i in persons_docs]
    persons_dict = {i["name"]: i for i in persons_docs}
    persons_string = "\n".join(persons)

    # Get the person who is interrogated
    interrogated_person = interrogation["person"]
    if interrogation["person"] in persons:
        doc = persons_dict[interrogated_person]
        _from = doc["_id"]
        _from_name = doc["_key"]
    text = interrogation["text"]

    # Make the text into smaller chunks and go through them
    chunks = text_splitter.split_text(text)

    print(len(chunks), 'chunks in', interrogation["person"].upper())

    earlier_mentions = []
    earlier_mentions_string = ''

    if len(chunks) > 1:
        part_of_string = ' en del av'
    else:
        part_of_string = ''
    for chunk in chunks:
        if earlier_mentions != []:
            earlier_mentions_list_string = '\n'.join(list(set(earlier_mentions)))
            earlier_mentions_string = f"Tidigare i förhöret har personerna nedan nämnts på följande sätt:\n{earlier_mentions_list_string}\n"
        
        prompt = f''''
Texten nedan är{part_of_string} ett polisförhör, kolla på den: \n

TEXT:
"""{chunk}""" \n
{interrogation["person"]} förhörs. Jag är intresserad av om någon eller några personer ur listan på personer nedan nämns i själva förhöret.\n
LISTA PÅ PERSONER:
{persons_string}\n
Nämns någon eller några av personerna i listan i förhöret? Personen kan nämnas med sitt fulla namn, men oftast bara förnamn eller efternamn. Försök alltså förstå om en person som nämns med förnamn är en person i listan.
{earlier_mentions_string}
Svara med FULLSTÄNDIGA namn från listan och hur personen nämns i texten på formen "namn;hur personen nämns\n". 
Nedan är ett exempel för att du ska förstå hur du ska svara:

<EXEMPEL>
John Lundqvist;John
Karl Renström; Karl
</EXEMPEL>

Svara ENBART med personens fullständiga namn och hur det nämns. Var noga med att använda formen "namn;hur personen nämns\n". 
Svara INTE med något resonemang eller något annat, och enbart med personer som nämns i förhöret. Om ingen person från listan nämns, svara med None.
Personer som nämns:'''

        relations = llm.generate(prompt)

        for relation in relations.replace('*', '').split("\n"):
            if relation == "None":
                continue
            try:
                name, mention = relation.split(";", 1)
                name = name.strip()
                mention = mention.strip()
                
                # Remove numbers at the beginning of the name using regex
                name = re.sub(r"^\d+\.\s*", "", name)
                name = re.sub(r"[^a-zA-Z\s-]", "", name)

                original_name = name
                
            except ValueError:
                print("\033[91m" + relation + "\033[0m")
                continue

            if name in persons:
                doc = persons_dict[name]
                _to = doc["_id"]
                _to_name = doc["_key"]

                earlier_mentions.append(f'{name} - {mention}')
            else:

                name_parts = name.split(" ")
                part1 = name_parts[0].strip()
                part2 = name_parts[1].strip() if len(name_parts) > 1 else ""
                if f'{part2} {part1}' in persons:
                    doc = persons_dict[f'{part2} {part1}']
                    _to = doc["_id"]
                    _to_name = doc["_key"]
                else:
                    arango_doc = {
                        "name": name,
                        "_key": arango.fix_key_name(name),
                    }
                    # pprint(mention_context)
                    # add = input(f"Add {name} ({mention}) to database? (y/n) >> ")
                    # if add in ["y", ""]:
                    if arango.fix_key_name(name) not in arango.db.collection("other_persons"):
                        doc = arango.db.collection("other_persons").insert(arango_doc)
                        _to = doc["_id"]
                        _to_name = doc["_key"]
                    else:
                        doc = arango.db.collection("other_persons").get(arango_doc["_key"])
                        _to = doc["_id"]
                        _to_name = doc["_key"]
            if _from_name == _to_name:
                continue

            relation_key = arango.fix_key_name(f"{_to}__{interrogation['_key']}").replace("persons_", "")
            if arango.db.has_document("all_relations/" + relation_key):
                continue
            
            # Ask LLM about the context of the mention
            prompt = f'Nedan är en del av ett förhör med {interrogated_person}.\n\n"""{chunk}"""\n\n{interrogated_person} nämner i förhöret en person vid namn {original_name}. Exakt vad säger {interrogated_person} om {original_name}? Svara så kortfattat som möjligt.\n\nSvar:'
            mention_context = llm.generate(prompt)

            prompt = f'Det här är ett svar från en assistent: """{mention_context}""".\nVerkar {original_name} nämnas? Svara bara med TRUE eller FALSE.\n\nSvar:'
            is_mentioned = llm.generate(prompt)
            if 'false' in is_mentioned.lower():
                print("\033[91m" + f"{original_name} not mentioned" + "\033[0m")
                print(mention_context, '\n')
                continue
            arango_doc = {
                    "_key": relation_key,
                    "_from": _from,
                    "_to": _to,
                    "_in": interrogation["_id"],
                    "context": "interrogation",
                    "mentioned_as": mention,
                    "mention": mention_context,
                    "date": interrogation["date"],
                }
            
            print("\033[92m" + f'{_from} -> {_to}' + "\033[0m" + f' ({interrogation["_key"]})')
            print(mention_context)
            print()
            arango.db.collection("all_relations").insert(
                arango_doc,
                overwrite=True,
            )