# Contains only NLP zero-shot classification function codes!
import sqlite3
from flask import g
from transformers import pipeline

# Set up zero-shot classification pipeline
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# The main "breakdown" functionality. Performs NLP using zero-shot!
def breakdown(userStory):
    
    # The results are stored into a dictionary variable with predifined labels
    processedStory = classifier(userStory, candidate_labels=["developer", "tester", "project manager", "system admin"])
   
    # Extract labels and scores
    scores = processedStory['scores']
    labels = processedStory['labels']
    
    # As the index of a score is always equal to its associated label,
    # We only need to index of the score to find correct label.
    maxScoreIndex = scores.index(max(scores)) # Gets the index of the highest score/accuracy
    maxLabel = labels[maxScoreIndex] # Gets the associated label name
    
    # Return labels 
    return maxLabel

# Function to grab all contents in the "Breakdown" table (except for unique ids)
# If adding any additional attributes to the table, this has to be updated accordingly
def getBreakdownContents():
    db = getattr(g, '_database', None) # Gets the _database attribute from the 'g' object. If it does not exist, returns 'None'
    if db is None:
        db = g._database = sqlite3.connect('Refineverse.db') # If db is None, create a new connection for db and g._database.
        cursor = db.cursor() # Creates a cursor object to handle data
        cursor.execute("SELECT user_story, assignedLabel FROM Breakdown") # The cursor executes the query
        rows = cursor.fetchall() # Stores the results of fetchall() into a variable
    return rows

# Function to insert a new row into the "Breakdown" table
# Using "with" for the connection here seems important, as otherwise it results in an exception
def insertBreakdownRow(user_story, assigned_label):
    with sqlite3.connect('Refineverse.db') as conn: # 'With' will automatically take care of closing and opening the connection
        cursor = conn.cursor()
        cursor.execute("INSERT INTO Breakdown (user_story, assignedLabel) VALUES (?, ?)", (user_story, assigned_label))
        conn.commit()