import streamlit as st
import sympy as sp
from solver import solve_equation, generate_steps
from plotter import plot_function
from utils import load_css, initialize_session_state
from auth import (
    login_user, signup_user, is_logged_in, logout_user,
    update_profile, get_user_profile, change_password
)
from models import History, SessionLocal, init_db
import base64
from io import BytesIO

# Initialize database
init_db()

def save_to_history(equation: str, solution: str):
    """Save equation and solution to user's history."""
    if not is_logged_in():
        return

    db = SessionLocal()
    try:
        history = History(
            user_id=st.session_state.user_id,
            equation=equation,
            solution=solution
        )
        db.add(history)
        db.commit()
    finally:
        db.close()

def load_user_history():
    """Load user's solution history."""
    if not is_logged_in():
        return []

    db = SessionLocal()
    try:
        return db.query(History).filter(
            History.user_id == st.session_state.user_id
        ).order_by(History.created_at.desc()).all()
    finally:
        db.close()

def render_math_symbols():
    """Render mathematical symbols for input."""
    st.markdown("### Mathematical Operators")

    cols = st.columns(8)
    all_symbols = {
        # Basic arithmetic
        '×': '*',
        '÷': '/',
        '^': '^',
        '=': '=',
        '(': '(',
        ')': ')',

        # Functions and special operators
        '√': '√',
        '∫': '∫',
        'd/dx': 'd/dx',
        '!': '!',
        'ℒ': 'L',
        '∑': 'sum',
        '∏': 'prod',
        '|': '|',

        # Constants
        'π': 'pi',
        'e': 'e',
        'i': 'i',
        '∞': 'oo',

        # Trigonometric
        'sin': 'sin(',
        'cos': 'cos(',
        'tan': 'tan(',
        'csc': 'csc(',
        'sec': 'sec(',
        'cot': 'cot(',

        # Inverse trigonometric
        'sin⁻¹': 'sin⁻¹(',
        'cos⁻¹': 'cos⁻¹(',
        'tan⁻¹': 'tan⁻¹(',

        # Other functions
        'ln': 'ln(',
        'log': 'log(',
        'e^': 'e^',
        '|x|': 'abs(',
    }

    for i, (label, symbol) in enumerate(all_symbols.items()):
        col_idx = i % 8
        if cols[col_idx].button(label, key=f"btn_{label}", help=f"Insert {label}"):
            if 'equation' not in st.session_state:
                st.session_state.equation = ''
            st.session_state.equation += symbol

def render_profile_settings():
    """Render user profile settings page."""
    st.title("Profile Settings")

    user = get_user_profile(st.session_state.user_id)

    # Profile photo upload
    st.subheader("Profile Photo")
    uploaded_file = st.file_uploader("Choose a profile photo", type=['jpg', 'jpeg', 'png'])
    if uploaded_file:
        # Convert to base64
        bytes_data = uploaded_file.getvalue()
        base64_image = base64.b64encode(bytes_data).decode()
        success, message = update_profile(st.session_state.user_id, profile_photo=base64_image)
        if success:
            st.success("Profile photo updated!")
        else:
            st.error(message)

    # Display current photo if exists
    if user.profile_photo:
        st.image(base64.b64decode(user.profile_photo), width=150)

    # Personal Information
    st.subheader("Personal Information")
    with st.form("profile_form"):
        full_name = st.text_input("Full Name", value=user.full_name or "")
        email = st.text_input("Email", value=user.email or "")
        school = st.text_input("School", value=user.school or "")
        grade = st.text_input("Grade/Year", value=user.grade or "")

        if st.form_submit_button("Update Profile"):
            success, message = update_profile(
                st.session_state.user_id,
                full_name=full_name,
                email=email,
                school=school,
                grade=grade
            )
            if success:
                st.success("Profile updated successfully!")
            else:
                st.error(message)

    # Password Change
    st.subheader("Change Password")
    with st.form("password_form"):
        current_password = st.text_input("Current Password", type="password")
        new_password = st.text_input("New Password", type="password")
        confirm_password = st.text_input("Confirm New Password", type="password")

        if st.form_submit_button("Change Password"):
            if new_password != confirm_password:
                st.error("New passwords do not match")
            else:
                success, message = change_password(
                    st.session_state.user_id,
                    current_password,
                    new_password
                )
                if success:
                    st.success("Password updated successfully!")
                else:
                    st.error(message)

def render_auth_page():
    """Render login/signup page."""
    st.title("Mathematical Problem Solver")

    tab1, tab2 = st.tabs(["Login", "Sign Up"])

    with tab1:
        st.header("Login")
        login_username = st.text_input("Username", key="login_username")
        login_password = st.text_input("Password", type="password", key="login_password")

        if st.button("Login"):
            if login_user(login_username, login_password):
                st.success("Successfully logged in!")
                st.rerun()
            else:
                st.error("Invalid username or password")

    with tab2:
        st.header("Sign Up")
        signup_username = st.text_input("Username", key="signup_username")
        signup_password = st.text_input("Password", type="password", key="signup_password")
        confirm_password = st.text_input("Confirm Password", type="password")

        st.info("Password must be at least 8 characters long and contain at least one uppercase letter.")

        if st.button("Sign Up"):
            if signup_password != confirm_password:
                st.error("Passwords do not match")
            else:
                success, message = signup_user(signup_username, signup_password)
                if success:
                    st.success("Account created successfully!")
                    st.rerun()
                else:
                    st.error(message)

def main():
    # Initialize session state and load CSS
    initialize_session_state()
    load_css()

    if not is_logged_in():
        render_auth_page()
        return

    # Main app header with settings and logout buttons
    col1, col2, col3 = st.columns([6, 1, 1])
    with col1:
        st.title("Mathematical Problem Solver")
    with col2:
        if st.button("Settings"):
            st.session_state.show_settings = True
            st.rerun()
    with col3:
        if st.button("Logout"):
            logout_user()
            st.rerun()

    st.markdown(f"Welcome, {st.session_state.username}!")

    # Show settings or main app
    if st.session_state.get('show_settings', False):
        render_profile_settings()
        if st.button("Back to Calculator"):
            st.session_state.show_settings = False
            st.rerun()
    else:
        # Main calculator interface
        # Sidebar for history
        with st.sidebar:
            st.header("Solution History")
            history = load_user_history()
            if history:
                for idx, item in enumerate(history):
                    with st.expander(f"Problem {idx + 1}"):
                        st.write(f"Input: {item.equation}")
                        st.write(f"Solution: {item.solution}")
                        st.write(f"Date: {item.created_at.strftime('%Y-%m-%d %H:%M')}")

        # Input method selection
        input_method = st.radio(
            "Choose input method:",
            ["Type equation/expression", "Use camera", "Example problems"]
        )

        if input_method == "Type equation/expression":
            # Add math symbols
            render_math_symbols()

            # Equation input
            if 'equation' not in st.session_state:
                st.session_state.equation = ''

            equation = st.text_input(
                "Enter your equation or expression:",
                value=st.session_state.equation,
                help="""Examples:
                - Equation: x^2 + 2x + 1 = 0
                - Integration: ∫sin x
                - Derivative: d/dx(x^2)
                - Factorial: 5!
                - Laplace: ℒ(t^2)"""
            )
            st.session_state.equation = equation

        elif input_method == "Use camera":
            st.info("📸 Camera input feature is coming soon! For now, please type your equation or choose from examples.")
            equation = ""

        else:
            equation = st.selectbox(
                "Select an example:",
                [
                    "x^2 + 2x + 1 = 0",
                    "∫sin x",
                    "d/dx(x^3)",
                    "sin^2 x + cos^2 x",
                    "5!",
                    "2x + 3 = 7",
                    "e^x + 1 = 0",
                    "log(x) = 1"
                ]
            )

        if st.button("Solve"):
            if equation:
                try:
                    # Solve the equation or expression
                    solution = solve_equation(equation)
                    steps = generate_steps(equation)

                    # Display solution
                    st.markdown("### Solution")
                    st.write(solution)

                    # Display steps
                    st.markdown("### Step-by-step Solution")
                    for step in steps:
                        st.write(step)

                    # Plot the function if possible
                    try:
                        st.markdown("### Function Visualization")
                        fig = plot_function(equation)
                        st.plotly_chart(fig)
                    except Exception as plot_error:
                        st.info("Visualization not available for this type of expression.")

                    # Save to history
                    save_to_history(equation, solution)

                except Exception as e:
                    st.error(f"Error: {str(e)}")
            else:
                st.warning("Please enter an equation or expression")

if __name__ == "__main__":
    main()