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


class Section:
    def __init__(self, type, page, filename="Förhörsprotokoll.pdf"):
        self.type = type
        self.filename = filename
        self.text = ""
        self.person = ""
        self.role = ""
        self.topic = ""
        self.date = ""
        self.start_page = page
        self.key = ""
        self.llm = LLM(chat=True)

    def add_to_arango(self):

        # prompt = f'Sammanfatta texten nedan: \n\n """{self.text[:7000]}""" \n\nSammanfattning:'
        # llm = LLM(chat=False)
        # self.summary = llm.generate(prompt)
        # print_green(self.summary)

        arango_doc = {
                "_key": self.key,
                "person": self.person,
                "role": self.role,
                "topic": self.topic,
                "date": self.date,
                "page": self.start_page,
                "text": self.text,
                "filename": self.filename,
                #"summary": self.summary,
                "name": self.name,
            }
        arango.db.collection(self.type).insert(arango_doc, overwrite_mode="update")
        print(f"Added {self.type} to ArangoDB with key {self.key}")


    def extract_interrogation(self, text):
        self.person = self.llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n Vem är förhörd? Namnet står ofta på formen "Efternamn, Förnamn". Jag vill att su ska svara på formen "Förnamn Efternamn"'
        )
        number_of_names = len(self.person.split(' '))
        if number_of_names != 2:
            first_name = self.llm.generate(
                f'Personens formella namn är alltså {self.person}. Om du kollar på själva föhörstexten, är personens förnamn (det som personen kallas för)? Svara bara med ett förnamn.'
            )
            last_name = self.llm.generate(
                f'Och i efternamn?'
            )
            self.name = f'{first_name} {last_name}'
            print_rainbow(self.person, '->', self.name)
        else:
            self.name = self.person

        self.role = self.llm.generate(
            f'Vem är {self.person}? Svara så kort som möjligt med titel eller beskrivning.'
        )
        self.topic = self.llm.generate(
            f'Vad handlade förhöret om? Svara så kortfattat som möjligt.'
        )
        self.date = self.llm.generate(
            f'När ägde förhöret rum? Svara på formen YYY-MM-DD'
        )

        cursor = arango.db.aql.execute(f'for doc in interrogations filter doc.page == {self.start_page} return doc._key', count=True)
        if cursor.count() == 1:
            self.key = cursor.next()
        else:
            print_red("Could not find key")
            self.key = arango.fix_key_name(f"{self.person}_{self.date}_p.{self.start_page}")


    def extract_pm(self, text):
        self.person = self.llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n Vem är uppgiftslämnare? Svara på formen "Förnamn Efternamn" \n\nPM:'
        )
        self.role = self.llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n Vem är {self.person}? Svara "None" om det inte framgår. \n\nTitel på person:'
        )   
        self.topic = self.llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n Vad handlade informationen om? Svara så kortfattat som möjligt. Svara "None" om det inte framgår. \n\Svar:'
        )
        self.date = self.llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n När lämnades informationen? Svara på formen YYY-MM-DD \n\nDatum:'
        )
        self.key = arango.fix_key_name(f"{self.person}_{self.date}_p.{self.start_page}")



def new_interrogation(page, section):

    if section.text != "":
        section.add_to_arango()
    section = Section("interrogations", page.number)
    section.extract_interrogation(page.get_text())
    return section


def new_pm(page, section):
    if section.text != "":
        section.add_to_arango()
    section = Section("pms", page.number)
    section.extract_interrogation(page.get_text())
    return section

def is_new_interrogation(page_text, page):
    # Check if there is a new interrogation
    control_words_interrogation = [
        "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_interrogation = 0
    for word in control_words_interrogation:
        if word in page_text:
            n_control_words_interrogation += 1
    if n_control_words_interrogation >= 2:
        print_purple('New interrogation', page)
        return True

def is_new_pm(page_text, page):
    control_words_pm = [
            "PM",
            "Uppgiften avser",
            "Upprättad av",
            "Sätt på vilket uppgift lämnats",
            "Uppgiftslämnare",
        ]
    n_control_words_pm = 0
    for word in control_words_pm:
        if word in page_text:
            n_control_words_pm += 1
    if n_control_words_pm >= 2:
        print_blue('New PM', page)
        return True
    


# * Llama

# Open the PDF file
filename = "Förhörsprotokoll.pdf"
area = fitz.Rect(0, 40, 520, 800)  # To exlude the header

doc = fitz.open(f"/home/lasse/mala/Förhörsprotokoll.pdf")

section = Section("interrogations", 0)


for page in doc.pages(9, len(doc) - 1):

    # Get the text from the page
    page_text = page.get_text("text")

    if is_new_interrogation(page_text, page):
        section = new_interrogation(page, section)
        area = fitz.Rect(0, 400, 520, 800)
    
    elif is_new_pm(page_text, page):
        # Check if there is a new PM
        area = fitz.Rect(0, 400, 520, 800)
        section = new_pm(page, section)
    
    else:
        # It's a "normal" page
        area = fitz.Rect(0, 40, 520, 800)  # To exlude the header

    blocks = page.get_text("blocks", clip=area)
    for block in blocks:
        section.text += block[4] + "\n\n"

new_interrogation(page, section)