import os
import streamlit as st
import requests
import json
from docx import Document
from io import BytesIO

# --- Set Streamlit environment variables for HuggingFace compatibility ---
# This helps prevent permission errors related to Streamlit's internal file operations
# and ensures it doesn't try to write to restricted directories.
os.environ["STREAMLIT_BROWSER_GATHER_USAGE_STATS"] = "false"
os.environ["STREAMLIT_SERVER_ENABLE_ARROW_IPC"] = "false"
os.environ["STREAMLIT_SERVER_FOLDER"] = "/tmp" # Added this line to specify a writable folder

# --- Page Configuration ---
st.set_page_config(
    page_title="Music Lesson Planner",
    page_icon="🎶",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- Constants and API Setup ---
# IMPORTANT: Set your Google API Key as an environment variable named GOOGLE_API_KEY
# You can get one from Google AI Studio: https://aistudio.google.com/app/apikey
GEMINI_API_KEY = os.getenv('GOOGLE_API_KEY')

# Base URL for Gemini API
GEMINI_API_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/models/"

# Available Gemini Models for comparison
GEMINI_MODELS = {
    "Gemini 2.0 Flash": "gemini-2.0-flash",
    "Gemini 1.5 Pro": "gemini-1.5-pro-latest",
    "Gemini 1.0 Pro": "gemini-1.0-pro",
}


# --- Helper Function for LLM API Call ---
def call_gemini_api(model_name, prompt_text, response_schema=None):
    """
    Calls the Gemini API with the given model and prompt.
    Handles JSON parsing and error reporting.
    """
    if not GEMINI_API_KEY:
        st.error(
            "Gemini API Key is not set. Please set the GOOGLE_API_KEY environment variable or replace `GEMINI_API_KEY = os.getenv('GOOGLE_API_KEY')` with your actual API key.")
        return None

    model_id = GEMINI_MODELS.get(model_name)
    if not model_id:
        st.error(f"Unknown model: {model_name}")
        return None

    url = f"{GEMINI_API_BASE_URL}{model_id}:generateContent?key={GEMINI_API_KEY}"

    headers = {
        "Content-Type": "application/json",
    }

    payload = {
        "contents": [
            {
                "role": "user",
                "parts": [{"text": prompt_text}]
            }
        ],
        "generationConfig": {}  # Initialize generationConfig
    }

    # If a response_schema is provided, configure for structured output
    if response_schema:
        payload["generationConfig"]["responseMimeType"] = "application/json"
        payload["generationConfig"]["responseSchema"] = response_schema
        # Add a clear instruction to the prompt to ensure JSON output
        # This can sometimes help guide the model more effectively.
        prompt_text = f"{prompt_text}\n\nPlease provide the response strictly in JSON format according to the schema provided, with no additional text or markdown outside the JSON object."
        payload["contents"][0]["parts"][0]["text"] = prompt_text

    try:
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()  # Raise an exception for HTTP errors (4xx or 5xx)
        response_data = response.json()

        if response_data and response_data.get("candidates"):
            # Access the text part of the response
            if response_schema:
                # For structured responses, the content is directly the JSON string
                raw_json_text = response_data["candidates"][0]["content"]["parts"][0]["text"]
                try:
                    # Attempt to parse the JSON string
                    parsed_json = json.loads(raw_json_text)
                    return parsed_json
                except json.JSONDecodeError as e:
                    st.error(f"Failed to parse JSON for {model_name} outline. Error: {e}")
                    st.text_area(f"Raw output from {model_name}:", raw_json_text, height=200)
                    return None
            else:
                # For unstructured responses, return the text directly
                return response_data["candidates"][0]["content"]["parts"][0]["text"]
        else:
            st.error(f"No valid response candidates found for {model_name}.")
            st.json(response_data)  # Display the full response for debugging
            return None

    except requests.exceptions.HTTPError as http_err:
        st.error(f"HTTP error occurred for {model_name}: {http_err}")
        st.error(f"Response content: {response.text}")
        return None
    except requests.exceptions.ConnectionError as conn_err:
        st.error(f"Connection error occurred for {model_name}: {conn_err}")
        return None
    except requests.exceptions.Timeout as timeout_err:
        st.error(f"Timeout error occurred for {model_name}: {timeout_err}")
        return None
    except requests.exceptions.RequestException as req_err:
        st.error(f"An unexpected error occurred for {model_name}: {req_err}")
        return None
    except Exception as e:
        st.error(f"An unexpected error occurred during API call for {model_name}: {e}")
        return None


# --- DOCX Conversion Function ---
def create_docx_file(outline_data, draft_text, lesson_topic, lesson_length, model_name):
    """
    Creates a DOCX file from outline data and draft text.
    """
    document = Document()

    # Add Lesson Title
    document.add_heading(outline_data.get("lessonTitle", "Untitled Lesson"), level=1)
    document.add_paragraph(f"Topic: {lesson_topic}")
    document.add_paragraph(f"Length: {lesson_length}")
    document.add_paragraph(f"Generated by: {model_name}")
    document.add_paragraph("\n")  # Add a blank line for spacing

    # Add Learning Objectives
    document.add_heading("Learning Objectives", level=2)
    for obj in outline_data.get("learningObjectives", []):
        document.add_paragraph(obj, style='List Bullet')
    document.add_paragraph("\n")

    # Add Materials
    document.add_heading("Materials", level=2)
    for material in outline_data.get("materials", []):
        document.add_paragraph(material, style='List Bullet')
    document.add_paragraph("\n")

    # Add Procedure
    document.add_heading("Procedure", level=2)
    for section in outline_data.get("procedure", []):
        document.add_heading(f"{section.get('sectionTitle', 'Section')} ({section.get('timeAllocation', 'N/A')})",
                             level=3)
        for activity in section.get("activities", []):
            document.add_paragraph(activity, style='List Bullet')
        document.add_paragraph("\n")

    # Add Assessment
    document.add_heading("Assessment", level=2)
    for assessment_item in outline_data.get("assessment", []):
        document.add_paragraph(assessment_item, style='List Bullet')
    document.add_paragraph("\n")

    # Add Full Draft Content
    document.add_heading("Full Lesson Draft", level=2)
    document.add_paragraph(draft_text)

    # Save document to a BytesIO object
    byte_io = BytesIO()
    document.save(byte_io)
    byte_io.seek(0)  # Rewind the buffer to the beginning
    return byte_io.getvalue()


# --- Password Protection ---
def authenticate_user():
    st.markdown("## 🔐 Secure Login")
    password = st.text_input("Password", type="password", key="password_input")
    submit = st.button("Login", key="login_button")

    if submit:
        correct_password = os.getenv("APP_PASSWORD")

        if password == correct_password:
            st.session_state["authenticated"] = True
            st.rerun()  # Rerun to clear password input and show app content
        else:
            st.error("Invalid password")


# Check authentication status
if "authenticated" not in st.session_state:
    st.session_state["authenticated"] = False

if not st.session_state["authenticated"]:
    authenticate_user()
    st.stop()  # Stop execution if not authenticated
else:
    # --- Main Application Logic (Protected by Authentication) ---
    st.title("🎶 Music Lesson Planner")

    st.markdown("""
    This app helps you draft outlines and detailed lesson plans for online music lessons using different Gemini models.
    Compare outputs to find the best fit for your pedagogical needs!
    """)

    # Initialize session state for outlines and drafts if not already present
    if 'outlines' not in st.session_state:
        st.session_state.outlines = {}
    if 'drafts' not in st.session_state:
        st.session_state.drafts = {}

    # Input fields for lesson details
    with st.sidebar:
        st.header("Lesson Details")
        lesson_topic = st.text_input("Lesson Topic", "Introduction to Solfege")
        lesson_length = st.selectbox("Lesson Length", ["2-minute", "5-minute", "10-minute", "15-minute"], index=1)

        st.header("Model Selection")
        selected_models = st.multiselect(
            "Select Gemini Models",
            list(GEMINI_MODELS.keys()),
            default=["Gemini 2.0 Flash", "Gemini 1.5 Pro"]
        )

        st.header("Prompt Customization")
        default_outline_system_prompt = (
            "You are an AI assistant specialized in creating concise and structured outlines for online music lessons. "
            "Your goal is to provide a clear, pedagogical framework that music educators can easily follow. "
            "Focus on key components: lesson title, learning objectives, materials, a step-by-step procedure with time allocations and activities, and assessment methods."
            "The lessons are online and asynchronous, so ensure the outline is suitable for self-paced learning."
        )
        outline_system_prompt = st.text_area("Outline System Prompt", default_outline_system_prompt, height=150)

        default_outline_user_prompt_template = (
            "Create a {lesson_length} online music lesson outline on the topic of '{lesson_topic}'. "
            "The outline should be structured as a JSON object with the following keys: "
            "'lessonTitle', 'learningObjectives' (list of strings), 'materials' (list of strings), "
            "'procedure' (list of objects, each with 'sectionTitle', 'timeAllocation', and 'activities' (list of strings)), "
            "and 'assessment' (list of strings). "
            "Ensure 'timeAllocation' for each procedure section is a string indicating duration (e.g., '5 minutes'). "
            "Make sure the total time allocations add up to the {lesson_length}."
            "Example for 'procedure' section: "
            '{{"sectionTitle": "Introduction", "timeAllocation": "5 minutes", "activities": ["Greet students", "Review previous concepts"]}}'
        )
        outline_user_prompt_template = st.text_area("Outline User Prompt Template",
                                                    default_outline_user_prompt_template, height=250)

        default_draft_system_prompt = (
            "You are an AI assistant specialized in expanding structured lesson outlines into detailed, engaging rough drafts for online music lessons. "
            "Your goal is to provide specific examples, pedagogical details, and interactive elements for each activity. "
            "The language should be engaging and professional, tailored for music educators."
            "The lessons are online and asynchronous, so ensure the draft is suitable for self-paced learning."
        )
        draft_system_prompt = st.text_area("Draft System Prompt", default_draft_system_prompt, height=150)

        generate_button = st.button("Generate Lesson Plans")

        # Add a logout button to the sidebar
        if st.session_state["authenticated"]:
            if st.button("Logout", key="logout_button_sidebar"):
                st.session_state["authenticated"] = False
                st.rerun()

    # --- Define Outline Schema ---
    outline_response_schema = {
        "type": "OBJECT",
        "properties": {
            "lessonTitle": {"type": "STRING"},
            "learningObjectives": {"type": "ARRAY", "items": {"type": "STRING"}},
            "materials": {"type": "ARRAY", "items": {"type": "STRING"}},
            "procedure": {
                "type": "ARRAY",
                "items": {
                    "type": "OBJECT",
                    "properties": {
                        "sectionTitle": {"type": "STRING"},
                        "timeAllocation": {"type": "STRING"},
                        "activities": {"type": "ARRAY", "items": {"type": "STRING"}}
                    },
                    "required": ["sectionTitle", "timeAllocation", "activities"]
                }
            },
            "assessment": {"type": "ARRAY", "items": {"type": "STRING"}}
        },
        "required": ["lessonTitle", "learningObjectives", "materials", "procedure", "assessment"]
    }

    # --- Lesson Generation Logic (Triggered by button) ---
    if generate_button:
        # Clear previous results when new generation is triggered
        st.session_state.outlines = {}
        st.session_state.drafts = {}
        st.session_state.lesson_topic = lesson_topic  # Store for download
        st.session_state.lesson_length = lesson_length  # Store for download
        st.session_state.selected_models = selected_models  # Store for download

        # Generate outlines
        for model_name in selected_models:
            current_outline_user_prompt = outline_user_prompt_template.format(
                lesson_length=lesson_length,
                lesson_topic=lesson_topic
            )
            full_outline_prompt = f"{outline_system_prompt}\n{current_outline_user_prompt}"
            outline_data = call_gemini_api(model_name, full_outline_prompt, outline_response_schema)
            if outline_data:
                st.session_state.outlines[model_name] = outline_data
            else:
                st.session_state.outlines[model_name] = None  # Mark as failed

        # Generate drafts
        for model_name in selected_models:
            if model_name in st.session_state.outlines and st.session_state.outlines[model_name]:
                outline_for_draft = json.dumps(st.session_state.outlines[model_name], indent=2)
                draft_prompt = (
                    f"{draft_system_prompt}\n\n"
                    f"Expand the following lesson outline into a detailed rough draft for a {lesson_length} lesson. "
                    "Provide specific examples and pedagogical details for each activity. "
                    "Ensure the language is engaging for music educators.\n\n"
                    f"Outline:\n```json\n{outline_for_draft}\n```"
                )
                raw_draft_text = call_gemini_api(model_name, draft_prompt)
                if raw_draft_text:
                    st.session_state.drafts[model_name] = raw_draft_text
                else:
                    st.session_state.drafts[model_name] = None  # Mark as failed
            else:
                st.session_state.drafts[model_name] = None  # Cannot generate draft without outline

    # --- Display Generated Content and Download Buttons (Always displayed if in session_state) ---
    if st.session_state.get('outlines') or st.session_state.get('drafts'):
        st.subheader("Generated Outlines")
        outline_cols = st.columns(len(st.session_state.get('selected_models', [])))
        for i, model_name in enumerate(st.session_state.get('selected_models', [])):
            with outline_cols[i]:
                st.markdown(f"### {model_name} Outline")
                if st.session_state.outlines.get(model_name):
                    st.json(st.session_state.outlines[model_name])
                else:
                    st.markdown(
                        f'<div class="result-box text-gray-500">Could not generate outline for {model_name}.</div>',
                        unsafe_allow_html=True)

        st.subheader("Generated Rough Drafts")
        draft_cols = st.columns(len(st.session_state.get('selected_models', [])))
        for i, model_name in enumerate(st.session_state.get('selected_models', [])):
            with draft_cols[i]:
                st.markdown(f"### {model_name} Rough Draft")
                if st.session_state.drafts.get(model_name):
                    st.markdown(f'<div class="result-box">{st.session_state.drafts[model_name]}</div>',
                                unsafe_allow_html=True)
                else:
                    st.markdown(
                        f'<div class="result-box text-gray-500">Could not generate draft for {model_name}.</div>',
                        unsafe_allow_html=True)

        st.subheader("Download Results")
        download_cols = st.columns(len(st.session_state.get('selected_models', [])))
        for i, model_name in enumerate(st.session_state.get('selected_models', [])):
            with download_cols[i]:
                outline_data = st.session_state.outlines.get(model_name)
                draft_text = st.session_state.drafts.get(model_name)

                if outline_data and draft_text:
                    docx_file = create_docx_file(
                        outline_data,
                        draft_text,
                        st.session_state.lesson_topic,
                        st.session_state.lesson_length,
                        model_name
                    )
                    st.download_button(
                        label=f"Download {model_name} (DOCX)",
                        data=docx_file,
                        file_name=f"{model_name.replace(' ', '_')}_lesson_plan_{st.session_state.lesson_topic.replace(' ', '_')}.docx",
                        mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                        key=f"download_docx_{model_name}"
                    )
                else:
                    st.markdown(
                        f'<div class="text-gray-500">Cannot download DOCX for {model_name} (missing data).</div>',
                        unsafe_allow_html=True)
    elif generate_button:
        st.info("No content generated. Please check for API errors or adjust prompts.")

    st.markdown("""
    <style>
    .result-box {
        border: 1px solid #ddd;
        border-radius: 8px;
        padding: 15px;
        margin-bottom: 10px;
        background-color: #f9f9f9;
        white-space: pre-wrap;
        word-wrap: break-word;
    }
    .text-gray-500 {
        color: #6b7280;
    }
    </style>
    """, unsafe_allow_html=True)