import streamlit as st
from azure.cosmos import CosmosClient
from dotenv import load_dotenv
import os

load_dotenv()

def handle_submit():
    client: CosmosClient = st.session_state["db"]
    database = client.get_database_client("assessment-database")
    container = database.get_container_client("assessments")
    assessment_item_coding_tasks = st.session_state["assessment_db_object"]["coding_tasks"]
    if st.session_state["upload_type"] == "Text upload":
        for i, coding_task in enumerate(st.session_state["assessment_db_object"]["coding_tasks"]):
            assessment_item_coding_tasks[i]["candidate_solution"] = st.session_state["code_text_"+str(i+1)]
    else:
        for i, coding_task in enumerate(st.session_state["assessment_db_object"]["coding_tasks"]):
            # print(st.session_state["file_upload_"+str(i+1)].read())
            try:
                assessment_item_coding_tasks[i]["candidate_solution"] = st.session_state["file_upload_"+str(i+1)].read().decode("utf-8")
            except:
                st.error("Please upload a valid code file")
                return
    upsert_item = st.session_state["assessment_db_object"]
    upsert_item["coding_tasks"] = assessment_item_coding_tasks
    upsert_item["process_status"] = "coding_solutions_uploaded"
    container.upsert_item(upsert_item)
    st.session_state["file_uploaded"] = True

if "db" not in st.session_state:
    client = CosmosClient(os.environ["COSMOS_DB_ENDPOINT"], os.environ["COSMOS_DB_KEY"])   
    st.session_state["db"] = client
if "assessment_db_object" not in st.session_state:
    if "assessment-id" in st.query_params:
        client: CosmosClient = st.session_state["db"]
        database = client.get_database_client("assessment-database")
        container = database.get_container_client("assessments")
        assessment = container.query_items(query="SELECT * FROM c WHERE c.id = @assessment_id", parameters=[dict(name="@assessment_id", value=st.query_params["assessment-id"]) ,], enable_cross_partition_query=True, max_item_count=1)
        st.session_state["assessment_db_object"] = list(assessment)[0]
    else:
        st.error("No assessment found. Please contact the side admin.", icon="🚨")
        st.stop()
if "file_uploaded" not in st.session_state:
    st.session_state["file_uploaded"] = False

col1, col2 = st.columns([2, 1])

col1.title("Code Upload")
col2.image("https://www.workgenius.com/wp-content/uploads/2023/03/WorkGenius_navy-1.svg")

if not st.session_state["file_uploaded"]:
    st.write("Assessment: "+st.session_state["assessment_db_object"]["assessment_title"])
    st.radio("Select text or file upload", ["Text upload", "File upload"], key="upload_type")

    if st.session_state["upload_type"] == "Text upload":
        for i, coding_task in enumerate(st.session_state["assessment_db_object"]["coding_tasks"]):
            st.write("Coding task "+str(i+1)+". : "+coding_task["coding_task"])
            st.text_area("Enter your code here", height=200, key="code_text_"+str(i+1))
    else:
        for i, coding_task in enumerate(st.session_state["assessment_db_object"]["coding_tasks"]):
            st.write("Coding task "+str(i+1)+". : "+coding_task["coding_task"])
            st.file_uploader("Upload your code file here", type=None, key="file_upload_"+str(i+1), accept_multiple_files=False)

    st.button("Submit", on_click=handle_submit)
    
else:    
    st.success("Your code has been successfully uploaded. You will be notified once the assessment is evaluated.", icon="🎉")