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

llm = LLM(chat=False)
docs = [i for i in arango.db.collection('interrogations').all()]
sorted_docs = sorted(docs, key=lambda x: x['date'])

filename = "Förhörsprotokoll.pdf"
pdf = fitz.open(f"pdfs/{filename}")

for doc in sorted_docs:

    pdf_page = pdf[doc['page']]
    text = pdf_page.get_text()

    print(doc['person'])
    prompt = f'Kolla på texten nedan: \n\n """{text}""" \n\n Den förhörda personen heter {doc["person"]} Varför förhörs {doc["person"]}? Om det har något att göra med {doc["person"]}s titel eller yrke, svara med det, annars med eventuell annan anledning. Om det inte finns någon speciell anledning eller titel, svara "None". \n\:'
    answer = llm.generate(prompt)
    doc['reason'] = answer
    print("\033[92m" + answer + "\033[0m")
    print()
    arango.db.collection('interrogations').update(doc)