import queue
import threading
from _arango import arango
from _llm import LLM


# Initialize and start the LLM object
llm = LLM(start=True)

# Add requests to the LLM object
q = 'for doc in interrogations filter doc.summary == null return doc'
docs = [i for i in arango.db.aql.execute(q)]
print(len(docs))
llm.stop()
exit()
for doc in docs:
    text = doc['text']
    prompt = f'Kolla på texten nedan: \n\n """{text}""" \n\n Sammanfatta förhöret med fokus på vad som sades, inte var det hölls eller annat formalia. Svara så kort som möjligt men var noga med detaljer som händelser som beskrivs, namn, datum och platser.\nKort sammanfattning:'
    llm.add_request(doc['_key'], prompt)

# Signal that all requests have been added
llm.finish_adding_requests()

# Get the results
llm.get_results()



# Initialize the LLM object
llm = LLM(chat=False, model="llama3:8b-instruct-q5_K_M")

# Create a queue for requests and a queue for results
request_queue = queue.Queue()
result_queue = queue.Queue()

# Create an event to signal when all requests have been added
all_requests_added_event = threading.Event()
all_results_processed_event = threading.Event()

# Start a separate thread that runs the generate_concurrent method
threading.Thread(target=llm.generate_concurrent, args=(request_queue, result_queue, all_requests_added_event, all_results_processed_event)).start()

# Add requests to the request queue

interrogations = arango.db.collection('interrogations').all()
for doc in interrogations:
    text = doc['text']
    prompt = f'Kolla på texten nedan: \n\n """{text}""" \n\n Sammanfatta förhöret med fokus på vad som sades, inte var det hölls eller annat formalia. Svara så kort som möjligt men var noga med detaljer som händelser som beskrivs, namn, datum och platser.\nKort sammanfattning:'
    request_queue.put((doc['_key'], prompt))

# Signal that all requests have been added
all_requests_added_event.set()

# Process the results
while True:
    try:
        # Take a result from the result queue
        doc_id, summary = result_queue.get(timeout=1)
        print("\033[92m" + doc_id + "\033[0m", summary)
        # Update the document with the summary
        arango.db.collection('interrogations').update_match({'_key': doc_id}, {'summary': summary})
    except queue.Empty:
        # If the result queue is empty and all results have been processed, break the loop
        if all_results_processed_event.is_set():
            break
        else:
            continue