from _llm import LLM
import fitz
from _arango import arango
from openai import OpenAI
from pprint import pprint


def extract_interrogation(text):

    interrogated = llm.generate(
        f'Kolla på texten nedan: \n\n """{text}""" \n\n Vem är förhörd? Svara på formen "Förnamn Efternamn" \n\nFörhörd person:'
    )
    interrogated_role = llm.generate(
        f'Kolla på texten nedan: \n\n """{text}""" \n\n Vem är {interrogated}? \n\nTitel på förhörd person:'
    )
    interrogation_topic = llm.generate(
        f'Kolla på texten nedan: \n\n """{text}""" \n\n Vad handlade förhöret om? Svara så kortfattat som möjligt. \n\nFörhörets syfte:'
    )
    interrogation_date = llm.generate(
        f'Kolla på texten nedan: \n\n """{text}""" \n\n När ägde förhöret rum? Svara på formen YYY-MM-DD \n\nFörhörsdatum:'
    )

    print(f"Förhörd: {interrogated}")
    print(f"Förhörd roll: {interrogated_role}")
    print(f"Förhörets syfte: {interrogation_topic}")
    print(f"Förhörsdatum: {interrogation_date}")

    if not arango.db.has_document(
        "interrogations/"
        + arango.fix_key_name(f"{interrogated}_{interrogation_date}_p.{page.number}")
    ):
        interrogation_key = arango.fix_key_name(
            f"{interrogated}_{interrogation_date}_p.{page.number}"
        )

        arango.db.collection("interrogations").insert(
            {
                "_key": interrogation_key,
                "interrogated": interrogated,
                "role": interrogated_role,
                "topic": interrogation_topic,
                "date": interrogation_date,
                "page": page.number,
                "text": text,
                "filename": filename,
            }
        )

    else:
        interrogation_key = arango.fix_key_name(
            f"{interrogated}_{interrogation_date}_p.{page.number}"
        )

    return (
        interrogation_key,
        interrogated,
        interrogated_role,
        interrogation_topic,
        interrogation_date,
    )


def extract_relations(text, interrogated_person):
    prompt = f'''Nedan är en del av ett förhör med {interrogated_person}. Jag vill veta vilka relationer som på något vis nämns i texten. Dessa kan vara mellan {interrogated_person} och någon annan, mellan två personer som {interrogated_person} berättar om eller mellan en person och en organisation/plats.
    Svara på formen "person1;person2;relation\n". Om det inte finns någon relation, svara med None.

    Nedan är ett exempel för att du ska förstå hur du ska svara:
    <EXEMPEL>
    Text: """En solig dag promenerade Anna längs med stadens livliga gator. Plötsligt stannade hon upp när hon såg en bekant gestalt längre fram. Med ett leende gick hon fram och hälsade på personen, och de inledde en trevlig konversation. Ju mer de pratade, desto fler minnen från barndomen väcktes till liv. Till slut insåg de att de faktiskt var gamla klasskompisar från högstadiet. Skratten ekade när de mindes tokiga stunder och gemensamma vänner. Det var en oväntad men glädjande återförening mitt i vardagens trummer."""
    Relationer: Anna;Peter;klasskompisar från högstadiet\n
    </EXEMPEL>
    
    Text: """{text}"""\n

    Svara ENBART med relationerna, inga förklaringar eller exempel eller något annat. Kom ihåg att svara på formen "person1;person2;relation\n".
    Relationer:'''

    return llm.generate(prompt)


# * OpenAI
OPENAI_KEY = "sk-proj-lDgKqh9eTLpbuSEaR69XT3BlbkFJsw0QkuXuZmf08mt9X76h"
client = OpenAI(
    # This is the default and can be omitted
    api_key=OPENAI_KEY,
)

# * Llama
llm = LLM(chat=False, model="llama3:8b-instruct-q5_K_M")

# To check if the interrogation has been found
interrogation = False

# Open the PDF file
filename = "Förhörsprotokoll.pdf"
doc = fitz.open(f"pdfs/{filename}")

for page in doc:
    text = page.get_text()

    control_words = [
        "Förhörsdatum",
        "Förhör påbörjat",
        "Förhör avslutat",
        "Förhörssätt",
        "Typ av förhör",
        "Förhörsvittne",
    ]
    n_control_words = 0
    for word in control_words:
        if word in text:
            n_control_words += 1

    if n_control_words >= 2:
        print("\n\n")
        interrogation = True
        (
            interrogation_key,
            interrogated,
            interrogated_role,
            interrogation_topic,
            interrogation_date,
        ) = extract_interrogation(text)

    if not interrogation:
        continue

    # Extract relations from the page
    relations = extract_relations(text, interrogated_person=interrogated)
    
    for i in relations.split("\n"):
        if i == "None":
            continue
        relation_parts = i.split(";")
        person1 = relation_parts[0]
        person2 = relation_parts[1]
        relation = " - ".join(relation_parts[2:])
        prompt = f'Kolla på texten nedan: \n\n """{text}""" \n\n Vilken del av texten beskriver relationen "{relation}" mellan {person1} och {person2}? Svara med den ORDAGRANNA texten. \n\nDel av text:'
        relation_text = llm.generate(prompt)
        arango_document = {
            "_key": arango.fix_key_name(
                f"{person1}_{person2}_{relation}_p.{page.number}"
            ),
            "person1": person1,
            "person2": person2,
            "relation": relation,
            "page": page.number,
            "text": relation_text,
            "filename": filename,
            "iterrogation": interrogation_key,
            "interrogated_person": interrogated,
            "interrogation_date": interrogation_date,
        }

        pprint(arango_document)
        arango.db.collection("relations").insert(
            arango_document, silent=True, overwrite=True
        )
        print("\n")

    # result = client.chat.completions.create(
    #     messages=[
    #         {
    #             "role": "user",
    #             "content": prompt,
    #         }
    #     ],
    #     model="gpt-4",
    # )
    # print(result)
    # answer = result.choices[0].message.content
    # print(answer)

    # print('\n\n')
    # print(f"\033[92m{answer}\033[0m")
    # exit()