import streamlit as st
import fitz
from fitz import Page, Document
from _llm import LLM
import re
from person_identifier import PersonFinder
from print_color import *

def set_name():
    st.session_state.name = st.session_state.names.pop(0)

def highlight_name_in_pdf(page: Page, name: str):
    # Search for the word in the page
    rectangles = page.search_for(name, quads=True)

    # Highlight the found words
    page.add_highlight_annot(rectangles)

    # Convert the page to a pixmap
    pixmap = page.get_pixmap(dpi=300)

    # Save the pixmap to a new image file
    pixmap.save(image_filename, 'png')


def show_image(filename):
    # Display the new image file in Streamlit
    st.image(filename)

def get_page(page_number):

    # Create a new document
    new_doc = fitz.open()
    new_doc.insert_pdf(doc, from_page=page_number, to_page=page_number)

    page = new_doc[0]
    page.set_cropbox(fitz.Rect(0, 100, 520, 800))
    return page

@st.cache_resource()
def get_extractor():
    return PersonFinder()

st.set_page_config(layout="wide")

filename = "Förhörsprotokoll.pdf"
image_filename = "highlighted.png"

page_number = 89

if 'doc' not in st.session_state:
    st.session_state.doc = fitz.open(filename)

doc = st.session_state.doc
page = get_page(page_number)




if 'all_names' not in st.session_state:
    st.session_state.all_names = {}

if 'names' not in st.session_state or st.session_state.names is None:
    person_extractor = PersonFinder(st.session_state.all_names)
    st.session_state.names = person_extractor.extract_names(person_extractor, page.get_text())
    st.session_state.all_names = person_extractor.names
    print_blue(st.session_state.names)
    print_purple(st.session_state.all_names)

names = st.session_state.names

if 'name' not in st.session_state:
    st.session_state.name = names.pop(0)
name = st.session_state.name

st.markdown(f'#### {name}')

highlight_name_in_pdf(page, name)

col1, col2 = st.columns([5,2])
with col1:
    show_image(image_filename)
with col2:
    next = st.button("Next", on_click=set_name)

if next:
    if len(names) == 0:
        st.session_state.names = None
    else:
        highlight_name_in_pdf(page, st.session_state.name)