import streamlit as st
import openai
from openai import OpenAI
import time
import gspread
from oauth2client.service_account import ServiceAccountCredentials
import PyPDF2
import io
from datetime import datetime
from PIL import Image

# Add some custom CSS to improve the layout
st.markdown("""
    <style>
    .stImage {
        text-align: center;
        display: block;
        margin-left: auto;
        margin-right: auto;
    }
    
    .stTitle {
        text-align: center;
        padding-bottom: 20px;
    }
    
    div[data-testid="stVerticalBlock"] > div:has(div.stButton) {
        text-align: center;
        padding: 10px 0;
    }
    </style>
    """, unsafe_allow_html=True)

# Constants
WORD_LIMIT = 11000
DAILY_API_LIMIT = 30  # Set your desired limit per user per day

# Set up OpenAI client
client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])

# Google Sheets setup
scope = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"]
creds = ServiceAccountCredentials.from_json_keyfile_name("genexam-2c8c645ecc0d.json", scope)
client_gs = gspread.authorize(creds)
sheet = client_gs.open("GeneXam user").sheet1

def check_user_in_sheet(username):
    """Check if user exists in sheet"""
    try:
        users_list = sheet.col_values(1)  # UserID column
        if username in users_list:
            return True
        return False
    except Exception as e:
        st.error(f"Error checking user: {str(e)}")
        return False

def get_user_stats(username):
    """Get user's current API usage statistics"""
    try:
        users_list = sheet.col_values(1)  # UserID column
        row_number = users_list.index(username) + 1
        
        daily_count = int(sheet.cell(row_number, 2).value)    # DailyAPICount
        total_count = int(sheet.cell(row_number, 3).value)    # TotalAPICount
        last_used = sheet.cell(row_number, 4).value           # LastUsedDate
        
        return {
            'daily_count': daily_count,
            'total_count': total_count,
            'last_used': last_used
        }
    except Exception as e:
        st.error(f"Error getting user stats: {str(e)}")
        return None

def update_api_usage(username):
    """Update both daily and total API usage counts"""
    try:
        users_list = sheet.col_values(1)
        row_number = users_list.index(username) + 1
        today = datetime.now().strftime('%Y-%m-%d')
        
        # Get current values
        stats = get_user_stats(username)
        if not stats:
            return False, "Error retrieving user statistics"
        
        # Reset daily count if it's a new day
        daily_count = stats['daily_count']
        if stats['last_used'] != today:
            daily_count = 0
        
        # Check daily limit
        if daily_count >= DAILY_API_LIMIT:
            return False, f"You have reached your daily limit of {DAILY_API_LIMIT} generations. Please try again tomorrow."
        
        # Update counts
        new_daily_count = daily_count + 1
        new_total_count = stats['total_count'] + 1
        
        # Update all values in sheet
        sheet.update_cell(row_number, 2, new_daily_count)  # Update DailyAPICount
        sheet.update_cell(row_number, 3, new_total_count)  # Update TotalAPICount
        sheet.update_cell(row_number, 4, today)            # Update LastUsedDate
        
        return True, None
        
    except Exception as e:
        return False, f"Error updating API usage: {str(e)}"

def extract_text_from_pdf(pdf_file):
    """Simple PDF text extraction with word limit check"""
    try:
        pdf_reader = PyPDF2.PdfReader(io.BytesIO(pdf_file.read()))
        text_content = ""
        for page in pdf_reader.pages:
            text_content += page.extract_text() + "\n"
        
        word_count = len(text_content.split())
        if word_count > WORD_LIMIT:
            return None, f"PDF content exceeds {WORD_LIMIT:,} words (contains {word_count:,} words). Please use a shorter document."
        return text_content, None
    except Exception as e:
        return None, f"Error processing PDF: {str(e)}"

def generate_questions_with_retry(username, knowledge_material, question_type, cognitive_level, extra_instructions, case_based, num_choices=None, max_retries=3):
    """Generate questions and update API usage"""
    
    # Check and update API usage before generating
    can_generate, error_message = update_api_usage(username)
    if not can_generate:
        st.error(error_message)
        return None
    
    # Adjust number of questions based on type
    if question_type == "Multiple Choice":
        num_questions = 3
        format_instructions = f"""
        For each multiple choice question:
        1. Present the question clearly
        2. Provide {num_choices} choices labeled with A, B, C{', D' if num_choices > 3 else ''}{', E' if num_choices > 4 else ''} after get new line from question
        3. After all questions, provide an ANSWER KEY section with:
           - The correct answer letter for each question
           - A brief explanation of why this is the correct answer
        """
    elif question_type == "Fill in the Blank":
        num_questions = 10
        format_instructions = """
        For each fill-in-the-blank question:
        1. Present the question with a clear blank space indicated by _____
        2. After all questions, provide an ANSWER KEY section with:
           - The correct answer for each blank
           - A brief explanation of why this answer is correct
           - Any alternative acceptable answers if applicable
        """
    elif question_type == "True/False":
        num_questions = 5
        format_instructions = """
        For each true/false question:
        1. Present the statement clearly
        2. After all questions, provide an ANSWER KEY section with:
           - Whether the statement is True or False
           - A detailed explanation of why the statement is true or false
           - The specific part of the source material that supports this answer
        """
    else:  # Open-ended
        num_questions = 3
        format_instructions = """
        For each open-ended question:
        1. Present the question clearly
        2. After all questions, provide an ANSWER KEY section with:
           - A structured scoring checklist of key points (minimum 3-5 points per question)
           - Each key point should be worth a specific number of marks
           - Total marks available for each question
           - Sample answer that would receive full marks
           - Common points that students might miss
        """

    # Base prompt
    prompt = f"""Generate {num_questions} {question_type.lower()} exam questions based on {cognitive_level.lower()} level from the following material:

{knowledge_material}

{format_instructions}

{extra_instructions}

Please format the output clearly with:
1. Questions section (numbered 1, 2, 3, etc.)
2. Answer Key section (clearly separated from questions)
3. Each answer should include explanation for better understanding

Make sure all questions and answers are directly related to the provided material."""

    # Modify prompt for case-based medical situations
    if case_based:
        prompt = f"""Generate {num_questions} {question_type.lower()} case-based medical exam questions based on {cognitive_level.lower()} level.

Use this material as the medical knowledge base:
{knowledge_material}

Each question should:
1. Start with a medical case scenario/patient presentation
2. Include relevant clinical details
3. Ask about diagnosis, treatment, or management
4. Be at {cognitive_level.lower()} cognitive level

{format_instructions}

{extra_instructions}

Please format the output with:
1. Cases and Questions (numbered 1, 2, 3, etc.)
2. Detailed Answer Key section including:
   - Correct answers
   - Clinical reasoning
   - Key diagnostic or treatment considerations
   - Common pitfalls to avoid"""

    retries = 0
    while retries < max_retries:
        try:
            response = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": "You are an expert exam question generator with deep knowledge in medical education. Create clear, well-structured questions with detailed answer keys and explanations."},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.7,
                max_tokens=3000  # Increased to accommodate answers and explanations
            )
            return response.choices[0].message.content
        except Exception as e:
            retries += 1
            st.warning(f"Attempt {retries} failed. Retrying... Error: {str(e)}")
            if retries == max_retries:
                st.error(f"Failed to generate questions after {max_retries} attempts. Error: {str(e)}")
                return None
            time.sleep(2)

# Main Streamlit interface
# Initialize session state variables
if 'login_step' not in st.session_state:
    st.session_state.login_step = 'username'
if 'username' not in st.session_state:
    st.session_state.username = None

# Login system
if st.session_state.username is None:
    # Center align the content
    col1, col2, col3 = st.columns([1,2,1])
    
    with col2:
        # Display logo
        st.image("GenExam.png", width=200)  # Assuming the image is saved as logo.png
        
        st.title("Login")
        username_input = st.text_input("Enter your username:")

        if st.session_state.login_step == 'username' and st.button("Login", use_container_width=True):
            if username_input:
                if check_user_in_sheet(username_input):
                    stats = get_user_stats(username_input)
                    if stats:
                        st.success(f"Welcome, {username_input}! 👋")
                        st.info(f"""
                        📊 Your API Usage Statistics:
                        - Today's Usage: {stats['daily_count']}/{DAILY_API_LIMIT} generations
                        - Total All-Time Usage: {stats['total_count']} generations
                        """)
                        st.session_state.login_step = 'enter_app'
                else:
                    st.warning("Username not found. Please try again.")
            else:
                st.warning("Please enter a valid username.")

        if st.session_state.login_step == 'enter_app':
            if st.button("🎯 Enter GeneXam Application", use_container_width=True):
                st.session_state.username = username_input
                st.rerun()

        # Show instructions
        if st.session_state.login_step == 'username':
            st.markdown("""
            ### How to Login:
            1. Enter your username and click 'Login' to verify your account
            2. After verification, click 'Enter GeneXam Application' to start using the system
            """)

else:
    # Main application code (ส่วนที่เหลือเหมือนเดิม)
    st.title(f"Welcome to GeneXam, {st.session_state.username}! 🎓")
    
    # Show current usage stats
    stats = get_user_stats(st.session_state.username)
    if stats:
        remaining = DAILY_API_LIMIT - stats['daily_count']
        st.info(f"""
        📊 Usage Statistics:
        - Daily Generations Remaining: {remaining}/{DAILY_API_LIMIT}
        - Total All-Time Generations: {stats['total_count']}
        """)
    
    # Create tabs for input methods
    tab1, tab2 = st.tabs(["Text Input", "PDF Upload"])
    
    with tab1:
        knowledge_material = st.text_area("Enter knowledge material to generate exam questions:")
        word_count = len(knowledge_material.split())
        if word_count > WORD_LIMIT:
            st.error(f"Text exceeds {WORD_LIMIT:,} words. Please shorten your content.")
    
    with tab2:
        st.info(f"Maximum content length: {WORD_LIMIT:,} words")
        uploaded_file = st.file_uploader("Upload a PDF file", type="pdf")
        
        if uploaded_file is not None:
            pdf_content, error = extract_text_from_pdf(uploaded_file)
            if error:
                st.error(error)
            else:
                st.success("PDF processed successfully!")
                knowledge_material = pdf_content

    # Question generation options
    col1, col2 = st.columns(2)
    
    with col1:
        question_type = st.selectbox(
            "Select question type:",
            ["Multiple Choice", "Fill in the Blank", "Open-ended", "True/False"]
        )
        
        if question_type == "Multiple Choice":
            num_choices = st.selectbox("Select number of choices:", [3, 4, 5])
        
        cognitive_level = st.selectbox(
            "Select cognitive level:",
            ["Recall", "Understanding", "Application", "Analysis", "Synthesis", "Evaluation"]
        )
    
    with col2:
        case_based = st.checkbox("Generate case-based medical exam questions")
        extra_instructions = st.text_area("Additional instructions (optional):")

    # Generate questions button
    if st.button("Generate Questions"):
        if 'knowledge_material' in locals() and knowledge_material.strip():
            with st.spinner("Generating questions..."):
                questions = generate_questions_with_retry(
                    st.session_state['username'],
                    knowledge_material,
                    question_type,
                    cognitive_level,
                    extra_instructions,
                    case_based,
                    num_choices if question_type == "Multiple Choice" else None
                )
                
                if questions:
                    st.write("### Generated Exam Questions:")
                    st.write(questions)
                    
                    # Update displayed stats after generation
                    new_stats = get_user_stats(st.session_state['username'])
                    if new_stats:
                        remaining = DAILY_API_LIMIT - new_stats['daily_count']
                        st.info(f"""
                        📊 Updated Usage Statistics:
                        - Daily Generations Remaining: {remaining}/{DAILY_API_LIMIT}
                        - Total All-Time Generations: {new_stats['total_count']}
                        """)
                    
                    # Download button
                    st.download_button(
                        label="Download Questions",
                        data=questions,
                        file_name='generated_questions.txt',
                        mime='text/plain'
                    )
        else:
            st.warning("Please enter knowledge material or upload a PDF file first.")