import streamlit as st
from groq_client import (
    extract_file_id,
    download_public_notebook,
    review_notebook_content,
)

st.set_page_config(page_title="Notebook Reviewer", layout="wide")

# Page title and description
st.title("📓 Jupyter Notebook Reviewer with Groq AI")
st.markdown(
    """
    Upload a Jupyter Notebook (`.ipynb`) file or provide a Google Drive link to receive **detailed feedback** on each code cell.  
    Each cell's context, code, and output are analyzed for accuracy, readability, and efficiency.
    """
)

# Layout: Sidebar for input options
with st.sidebar:
    st.header("📥 Upload or Provide Link")
    uploaded_file = st.file_uploader("Upload Notebook File", type=["ipynb"])
    google_drive_link = st.text_input("OR Enter Google Drive Link")

# Main review section
st.subheader("📑 Notebook Review")

if st.button("Review Notebook"):
    notebook_content = None

    # Process notebook input
    if uploaded_file:
        notebook_content = uploaded_file.getvalue().decode("utf-8")
        st.success("Notebook uploaded successfully!")
    elif google_drive_link:
        try:
            file_id = extract_file_id(google_drive_link)
            notebook_content = download_public_notebook(file_id)
            st.success("Notebook downloaded successfully from Google Drive!")
        except Exception as e:
            st.error(f"Error: {e}")
            st.stop()
    else:
        st.error("Please upload a file or provide a Google Drive link.")
        st.stop()

    # Analyze the notebook content
    try:
        feedback = review_notebook_content(notebook_content)
        st.success(
            "Notebook successfully reviewed! Scroll down to see detailed feedback."
        )

        # Display feedback for each cell
        for cell_index, cell_feedback in enumerate(feedback, start=1):
            with st.expander(f"Feedback for Cell {cell_index}", expanded=False):
                # Splitting feedback into sections
                cell_lines = cell_feedback.split("\n")
                markdown_lines = []
                code_block = []

                for line in cell_lines:
                    if line.strip().startswith("```"):  # Identify code block markers
                        code_block.append(line.strip("`"))
                    elif code_block:
                        code_block.append(line)
                    else:
                        markdown_lines.append(line)

                # Render text feedback
                st.markdown("\n".join(markdown_lines))

                # Render code if present
                if code_block:
                    st.code("\n".join(code_block), language="python")

    except Exception as e:
        st.error(f"An error occurred: {e}")