import gradio as gr
import os
import time
import pandas as pd
import numpy as np
from typing import List
from supabase import create_client
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()
client = OpenAI(api_key=os.environ.get("API_TOKEN", ""))
supabase = create_client(os.getenv("SUPABASE_URL", ""), os.getenv("SUPABASE_KEY", ""))

def get_embedding(text: str) -> List[float]:
    r = client.embeddings.create(model="text-embedding-3-small", input=text)
    return [float(x) for x in r.data[0].embedding]

def cosine_similarity(a: List[float], b: List[float]) -> float:
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

def sign_up_school(
    username, password, name, address_line1, city, state, zip_code,
    school_type, educational_level, enrollment, title_i_status,
    free_reduced_lunch_percentage, contact_email, additional_info
):
    try:
        auth_response = supabase.auth.sign_up({"email": username, "password": password})
        user_id = auth_response.user.id
        profile_text = f"""
        School: {name}
        Type: {school_type}
        Level: {educational_level}
        Title I: {title_i_status}
        Free/Reduced Lunch: {free_reduced_lunch_percentage}%
        Additional Info: {additional_info}
        """
        profile_embedding = get_embedding(profile_text)
        if len(profile_embedding) != 1536:
            raise ValueError(f"Expected embedding dimension 1536, got {len(profile_embedding)}")
        school_data = {
            "user_id": user_id,
            "name": name,
            "address_line1": address_line1,
            "city": city,
            "state": state,
            "zip_code": zip_code,
            "school_type": school_type,
            "educational_level": educational_level,
            "enrollment": enrollment,
            "title_i_status": title_i_status,
            "free_reduced_lunch_percentage": free_reduced_lunch_percentage,
            "contact_email": contact_email,
            "additional_info": additional_info,
            "embedding": profile_embedding
        }
        supabase.table('schools').insert(school_data).execute()
        return True, "School profile created successfully!", user_id
    except Exception as e:
        return False, f"Error: {str(e)}", None

def login(username, password):
    try:
        r = supabase.auth.sign_in_with_password({"email": username, "password": password})
        user_id = r.user.id
        school_response = supabase.table('schools').select("*").eq('user_id', user_id).execute()
        school_id = None
        if school_response.data and len(school_response.data) > 0:
            school_id = school_response.data[0]['user_id']
        return True, user_id, school_id
    except Exception as e:
        return False, f"Error: {str(e)}", None

def get_school_name(school_id):
    if not school_id:
        return "No school selected"
    try:
        r = supabase.table('schools').select("name").eq('user_id', school_id).execute()
        if r.data and len(r.data) > 0:
            return r.data[0]['name']
        return "Unknown school"
    except:
        return "Error loading school"

def search_grants(query, school_id=None):
    try:
        headers = ["Title", "Provider", "Due_date", "Description", "Funding_amount", "Requirements", "Eligibility", "Categories", "Score"]
        if not query or query.strip() == "":
            return pd.DataFrame(columns=headers)
        query_embedding = get_embedding(query)
        school_data = None
        if school_id and str(school_id).strip().lower() != 'none':
            try:
                sr = supabase.table('schools').select("*").eq('user_id', school_id).execute()
                if sr.data:
                    school_data = sr.data[0]
                    if school_data.get('embedding'):
                        try:
                            school_data['embedding'] = eval(school_data['embedding'])
                        except:
                            school_data['embedding'] = None
            except:
                pass
        grants_response = supabase.table('grants').select("*").execute()
        if not grants_response.data:
            return pd.DataFrame(columns=headers)
        grants = grants_response.data
        scored_grants = []
        for g in grants:
            try:
                if g.get('embedding'):
                    try:
                        grant_embedding = eval(g['embedding'])
                    except:
                        continue
                    q_sim = cosine_similarity(query_embedding, grant_embedding)
                    s_sim = 0
                    if school_data and school_data.get('embedding'):
                        s_sim = cosine_similarity(school_data['embedding'], grant_embedding)
                    context_score = 0
                    scored = 0.4 * q_sim + 0.3 * s_sim + 0.3 * context_score
                    if scored > 0.1:
                        scored_grants.append({
                            'Title': g.get('title', ''),
                            'Provider': g.get('provider', ''),
                            'Due_date': g.get('due_date', ''),
                            'Description': g.get('description', ''),
                            'Funding_amount': g.get('funding_amount', ''),
                            'Requirements': g.get('requirements', ''),
                            'Eligibility': g.get('eligibility', ''),
                            'Categories': g.get('categories', ''),
                            'Score': round(scored, 3)
                        })
            except:
                continue
        if not scored_grants:
            return pd.DataFrame(columns=headers)
        scored_grants.sort(key=lambda x: x['Score'], reverse=True)
        df = pd.DataFrame(scored_grants)
        return df
    except:
        return pd.DataFrame(columns=headers)

def chat_with_ai(message, history, grant_id=None, school_id=None):
    print(f"Chat with AI called with message: {message}, grant_id: {grant_id}, school_id: {school_id}")
    try:
        grant_data = None
        school_data = None
        school_profile = None
        if school_id and str(school_id).strip().lower() != 'none':
            sr = supabase.table('schools').select("*").eq('user_id', school_id).execute()
            if sr.data:
                school_data = sr.data[0]
                school_profile = f"""
                School Profile:
                Name: {school_data.get('name', '')}
                Type: {school_data.get('school_type', '')}
                Level: {school_data.get('educational_level', '')}
                Title I Status: {school_data.get('title_i_status', '')}
                Free/Reduced Lunch: {school_data.get('free_reduced_lunch_percentage', '')}%
                Enrollment: {school_data.get('enrollment', '')}
                Location: {school_data.get('city', '')}, {school_data.get('state', '')}
                Additional Info: {school_data.get('additional_info', '')}
                """
        if grant_id and str(grant_id).strip().lower() != 'none':
            grr = supabase.table('grants').select("*").eq('grant_id', grant_id).execute()
            if not grr.data:
                grr = supabase.table('grants').select("*").eq('grant_id', grant_id).execute()
            if grr.data:
                grant_data = grr.data[0]
        system_prompt = f"""
        School Profile: {school_profile if school_profile else 'No school profile available'}
        
        Grant: {grant_data['title'] if grant_data else 'No specific grant selected'}
        
        You are GrantGPT, a grant writing assistant. Begin by summarizing the selected grant:
        "You are applying for [Grant Name], which provides [Funding Amount] for [Purpose]. The deadline is [Date]."
        
        Ask if they have a project in mind or need suggestions. Guide users through sections:
        - Project Abstract/Summary
        - Statement of Need
        - Program Description
        - Budget & Funding
        - Evaluation Plan
        
        Keep responses focused on grant writing assistance. Redirect unrelated questions to appropriate resources.
        """
        messages = [{"role": "system", "content": system_prompt}]
        if history:
            for user_msg, assistant_msg in history:
                messages.append({"role": "user", "content": user_msg})
                messages.append({"role": "assistant", "content": assistant_msg})
        messages.append({"role": "user", "content": message})
        response = client.chat.completions.create(model="gpt-4", messages=messages, max_tokens=2000, temperature=0.7)
        full_message = response.choices[0].message.content
        if (grant_id and str(grant_id).strip().lower() != 'none' and school_id and str(school_id).strip().lower() != 'none'):
            try:
                app_response = supabase.table('applications').select("id").eq('grant_id', grant_id).eq('school_id', school_id).execute()
                if not app_response.data:
                    app_response = supabase.table('applications').insert({
                        'grant_id': grant_id,
                        'school_id': school_id,
                        'status': 'Draft'
                    }).execute()
                if app_response.data:
                    application_id = app_response.data[0]['id']
                    supabase.table('application_chat_history').insert([
                        {'application_id': application_id, 'message_type': 'user', 'content': message},
                        {'application_id': application_id, 'message_type': 'assistant', 'content': full_message}
                    ]).execute()
            except:
                pass
        return full_message
    except Exception as e:
        return f"Error: {str(e)}"

custom_css = """
body {
    margin: 0;
    background: linear-gradient(to bottom, #c2e9fb, #e2fcff);
    font-family: 'Helvetica Neue', sans-serif;
    min-height: 100vh;
}
.gradio-container {
    max-width: 100% !important;
    margin: auto;
    padding: 0;
    display: flex;
    flex-direction: column;
    justify-content: center;
    min-height: 100vh;
}
.modal-overlay {
    position: fixed;
    top: 0; 
    left: 0; 
    right: 0; 
    bottom: 0;
    background: rgba(0, 0, 0, 0.2);
    z-index: 999;
    display: flex;
    justify-content: center;
    align-items: center;
    min-height: 100vh;
}
.auth-modal {
    background: #fff;
    padding: 40px;
    border-radius: 10px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    z-index: 1000;
    width: 400px !important;
    max-width: 90% !important;
    height: auto !important;
    min-height: 300px !important;
    max-height: 400px !important;
    overflow-y: auto;
    margin: 0 auto;
    display: flex !important;
    flex-direction: column !important;
    justify-content: flex-start !important;
}
.landing-wrapper {
    display: flex;
    justify-content: center;
    align-items: center;
    min-height: 100vh;
    width: 100%;
    padding: 20px;
    box-sizing: border-box;
}
.landing-buttons {
    display: flex;
    flex-direction: column;
    gap: 15px;
    width: 100%;
    margin-top: 20px;
}
.landing-button {
    background-color: #8ecae6 !important;
    color: #2c3e50 !important;
    border: none !important;
    border-radius: 25px !important;
    padding: 12px 20px !important;
    font-size: 16px !important;
    font-weight: 500 !important;
    cursor: pointer !important;
    transition: background-color 0.3s ease !important;
    width: 100% !important;
    text-align: center !important;
}
.landing-button:hover {
    background-color: #76b5da !important;
}
@media (max-width: 768px) {
    .auth-modal {
        width: 90% !important;
        padding: 20px;
    }
}
.auth-modal h1 {
    color: #2c3e50 !important;
    margin: 0 0 20px 0 !important;
    text-align: center;
    flex-shrink: 0;
}
.auth-modal button {
    background-color: #ffffff;
    color: #2c3e50;
    border: 1px solid #ccc;
    border-radius: 25px;
    padding: 0.6rem 1.2rem;
    font-size: 1rem;
    margin: 10px;
    cursor: pointer;
}
.auth-modal button:hover {
    background-color: #f0f0f0;
}
.login-form, .signup-form {
    margin-top: 20px;
}
.gr-text-input input[type=text],
.gr-text-input input[type=password],
textarea {
    background-color: #fff !important;
    color: #333 !important;
    border: 1px solid #ccc !important;
    border-radius: 4px !important;
}
h1, h2, h3, h4, h5, h6 {
    color: #2c3e50;
}
.grant-card {
    background: #ffffff;
    border-radius: 10px;
    padding: 20px;
    margin-bottom: 15px;
    box-shadow: 0 2px 5px rgba(0,0,0,0.1);
    cursor: pointer;
}
.grant-card:hover {
    box-shadow: 0 4px 8px rgba(0,0,0,0.15);
}
.nav-bar {
    display: flex;
    justify-content: space-between;
    align-items: center;
    background: #fff;
    padding: 10px 20px;
    border-bottom: 1px solid #ccc;
}
.nav-bar button {
    background: none;
    border: none;
    cursor: pointer;
    font-weight: bold;
    color: #2c3e50;
}
.nav-bar button:hover {
    text-decoration: underline;
}
.grants-container {
    max-height: 600px;
    overflow-y: auto;
    padding: 20px;
    background: rgba(255, 255, 255, 0.5);
    border-radius: 10px;
    margin: 20px 0;
}
.grants-container::-webkit-scrollbar {
    width: 8px;
}
.grants-container::-webkit-scrollbar-track {
    background: #f1f1f1;
    border-radius: 4px;
}
.grants-container::-webkit-scrollbar-thumb {
    background: #888;
    border-radius: 4px;
}
.grants-container::-webkit-scrollbar-thumb:hover {
    background: #555;
}
.help-button-container {
    display: flex;
    justify-content: center;
    align-items: center;
    margin-top: 2rem;
    margin-bottom: 2rem;
}
.centered-button {
    min-width: 200px;
    background-color: #2c3e50;
    color: white;
    border: none;
    border-radius: 25px;
    padding: 1rem 2rem;
    font-size: 1.1rem;
    cursor: pointer;
    transition: background-color 0.3s ease;
}
.centered-button:hover {
    background-color: #34495e;
}
.grant-details {
    background: white;
    padding: 2rem;
    border-radius: 10px;
    box-shadow: 0 2px 5px rgba(0,0,0,0.1);
    color: #333;
    margin: 1rem;
    max-height: 70vh;
    overflow-y: auto;
}
.grant-details::-webkit-scrollbar {
    width: 8px;
}
.grant-details::-webkit-scrollbar-track {
    background: #f1f1f1;
    border-radius: 4px;
}
.grant-details::-webkit-scrollbar-thumb {
    background: #888;
    border-radius: 4px;
}
.grant-details::-webkit-scrollbar-thumb:hover {
    background: #555;
}
.grant-details h1, 
.grant-details h2, 
.grant-details h3, 
.grant-details h4, 
.grant-details h5, 
.grant-details h6 {
    color: #2c3e50;
}
.grant-details p {
    color: #333;
    line-height: 1.6;
}
.grant-details strong {
    color: #2c3e50;
}
.grant-details ul, 
.grant-details ol {
    color: #333;
}
.grant-details li {
    color: #333;
    line-height: 1.6;
}
.grant-details .loading {
    display: flex;
    justify-content: center;
    align-items: center;
    padding: 2rem;
}
.grant-details .loading::after {
    content: "";
    width: 40px;
    height: 40px;
    border: 4px solid #f3f3f3;
    border-top: 4px solid #2c3e50;
    border-radius: 50%;
    animation: spin 1s linear infinite;
}
@keyframes spin {
    0% { transform: rotate(0deg); }
    100% { transform: rotate(360deg); }
}
"""

def build_login_form():
    with gr.Column(elem_classes="login-form"):
        email = gr.Textbox(label="Email")
        password = gr.Textbox(label="Password", type="password")
        login_button = gr.Button("Log In")
        error = gr.Markdown(visible=False)
    return email, password, login_button, error

def build_signup_form():
    with gr.Column(elem_classes="signup-form"):
        email = gr.Textbox(label="Email")
        password = gr.Textbox(label="Password", type="password")
        name = gr.Textbox(label="School Name")
        with gr.Row():
            address_line1 = gr.Textbox(label="Address Line 1")
            city = gr.Textbox(label="City")
        with gr.Row():
            state = gr.Dropdown(["CA", "NY", "TX"], label="State")
            zip_code = gr.Textbox(label="ZIP Code")
        with gr.Row():
            school_type = gr.Dropdown(["Public", "Private", "Charter", "Nonprofit", "Other"], label="School Type")
            educational_level = gr.Dropdown(["K-12", "Higher Education", "Both"], label="Educational Level")
        with gr.Row():
            enrollment = gr.Number(label="Enrollment")
            title_i_status = gr.Checkbox(label="Title I Status")
            free_reduced_lunch_percentage = gr.Number(label="Free/Reduced Lunch %")
        contact_email = gr.Textbox(label="Contact Email")
        additional_info = gr.Textbox(label="Additional Info", lines=3)
        signup_button = gr.Button("Create Profile")
        error = gr.Markdown(visible=False)
    return (
        email, password, name, address_line1, city, state, zip_code,
        school_type, educational_level, enrollment, title_i_status,
        free_reduced_lunch_percentage, contact_email, additional_info,
        signup_button, error
    )

with gr.Blocks(css=custom_css) as demo:
    page_state = gr.State("landing")
    current_school_id = gr.State(None)
    selected_grant = gr.State(None)
    selected_grant_id = gr.State(None)
    is_authenticated = gr.State(False)

    def router(page):
        return {
            landing_page: gr.update(visible=(page == "landing")),
            login_page: gr.update(visible=(page == "login")),
            signup_page: gr.update(visible=(page == "signup")),
            home_page: gr.update(visible=(page == "home")),
            grant_details_page: gr.update(visible=(page == "grant_details")),
            chat_page: gr.update(visible=(page == "chat"))
        }

    with gr.Column(visible=True) as landing_page:
        with gr.Column(elem_classes="modal-overlay"):
            with gr.Column(elem_classes="auth-modal"):
                gr.Markdown("# Welcome to GrantRight!")
                with gr.Column(elem_classes="landing-buttons"):
                    create_account_btn = gr.Button("CREATE ACCOUNT", elem_classes="landing-button")
                    login_btn = gr.Button("LOG INTO EXISTING ACCOUNT", elem_classes="landing-button")

        def go_signup():
            return "signup"
        def go_login():
            return "login"

        create_account_btn.click(go_signup, None, page_state)
        login_btn.click(go_login, None, page_state)

    with gr.Column(visible=False) as login_page:
        with gr.Column(elem_classes="auth-modal"):
            gr.Markdown("## Log In")
            login_email, login_password, login_submit, login_error = build_login_form()
            def handle_login_click(email, pwd):
                success, result, school_id = login(email, pwd)
                if success:
                    if school_id:
                        return {
                            is_authenticated: True,
                            current_school_id: school_id,
                            page_state: "home",
                            login_error: gr.update(visible=False)
                        }
                    return {
                        is_authenticated: True,
                        current_school_id: None,
                        page_state: "home",
                        login_error: gr.update(visible=False)
                    }
                else:
                    return {login_error: gr.update(visible=True, value=result)}
            login_submit.click(
                handle_login_click,
                inputs=[login_email, login_password],
                outputs=[is_authenticated, current_school_id, page_state, login_error]
            )

    with gr.Column(visible=False) as signup_page:
        with gr.Column(elem_classes="auth-modal"):
            gr.Markdown("## Create School Profile")
            (
                su_email, su_password, su_name, su_address, su_city,
                su_state, su_zip, su_type, su_level, su_enroll,
                su_titlei, su_lunch, su_contact, su_info,
                su_submit, su_error
            ) = build_signup_form()

            def handle_signup_click(
                username, password, name, address_line1, city, state, zip_code,
                school_type, educational_level, enrollment, title_i_status,
                free_reduced_lunch_percentage, contact_email, additional_info
            ):
                success, msg, school_id = sign_up_school(
                    username, password, name, address_line1, city, state, zip_code,
                    school_type, educational_level, enrollment, title_i_status,
                    free_reduced_lunch_percentage, contact_email, additional_info
                )
                if success:
                    return {
                        is_authenticated: True,
                        current_school_id: school_id,
                        page_state: "home",
                        su_error: gr.update(visible=False)
                    }
                else:
                    return {su_error: gr.update(visible=True, value=msg)}

            su_submit.click(
                handle_signup_click,
                inputs=[
                    su_email, su_password, su_name, su_address, su_city, su_state,
                    su_zip, su_type, su_level, su_enroll, su_titlei, su_lunch,
                    su_contact, su_info
                ],
                outputs=[is_authenticated, current_school_id, page_state, su_error]
            )

    with gr.Column(visible=False) as home_page:
        with gr.Row(elem_classes="nav-bar"):
            nav_home_btn = gr.Button("Home")
            school_label = gr.Markdown("")
        search_bar = gr.Textbox(placeholder="Example: We're looking for STEM education grants to establish a robotics program...")
        search_button = gr.Button("Search")
        search_loading_indicator = gr.Markdown("", visible=False, elem_classes="loading")
        grants_container = gr.Column(elem_classes="grants-container")
        grant_dropdown = gr.Dropdown(
            choices=[],
            label="Select a Grant",
            visible=False
        )
        selected_grant_id = gr.State(None)
        view_grant_btn = gr.Button(
            "View Grant Details",
            visible=False
        )
        def update_school_label(sid):
            if sid:
                return get_school_name(sid)
            return "No School Selected"
        current_school_id.change(update_school_label, current_school_id, school_label)
        def do_search(q, sid):
            df = search_grants(q, sid)
            if df.empty:
                return {
                    grants_container: gr.update(value="<p>No grants found matching your search.</p>"),
                    grant_dropdown: gr.update(choices=[], visible=False),
                    view_grant_btn: gr.update(visible=False),
                    selected_grant_id: None
                }
            df = df.nlargest(30, 'Score')
            html_content = ""
            for _, row in df.iterrows():
                html_content += f"""
                <div class="grant-card">
                    <h3>{row['Title']}</h3>
                    <p>{row['Description'][:80]}...</p>
                    <p>Due: {row['Due_date']}</p>
                    <p>Amount: {row['Funding_amount']}</p>
                    <p>Score: {row['Score']}</p>
                </div>
                """
            grants_data = {}
            for _, row in df.iterrows():
                grants_data[row['Title']] = row['id'] if 'id' in row else None
            choices = [(row['Title'], row['Title']) for _, row in df.iterrows()]
            return {
                grants_container: gr.update(value=html_content),
                grant_dropdown: gr.update(choices=choices, visible=True),
                view_grant_btn: gr.update(visible=True),
                selected_grant_id: None
            }
        def handle_grant_selection(title, sid):
            if not title:
                return "home", None, None
            r = supabase.table('grants').select("grant_id").eq('title', title).execute()
            grant_id = None
            if r.data and len(r.data) > 0:
                grant_id = r.data[0]['grant_id']
                print(f"Selected grant ID: {grant_id} for title: {title}")
            else:
                print(f"No grant ID found for title: {title}")
            return "grant_details", title, grant_id
        search_button.click(
            lambda: gr.update(visible=True),
            None,
            search_loading_indicator,
            queue=False
        ).then(
            do_search,
            inputs=[search_bar, current_school_id],
            outputs=[grants_container, grant_dropdown, view_grant_btn, selected_grant_id]
        ).then(
            lambda: gr.update(visible=False),
            None,
            search_loading_indicator,
            queue=False
        )

        view_grant_btn.click(
            handle_grant_selection,
            inputs=[grant_dropdown, current_school_id],
            outputs=[page_state, selected_grant, selected_grant_id]
        )

    with gr.Column(visible=False) as grant_details_page:
        with gr.Row(elem_classes="nav-bar"):
            nav_home_btn2 = gr.Button("Home")
            school_label2 = gr.Markdown("")
        with gr.Column(elem_classes="grant-details"):
            grant_title = gr.Markdown("")
            grant_info = gr.Markdown("")
            loading_indicator = gr.Markdown("", visible=False)
            with gr.Row(elem_classes="help-button-container"):
                help_write_btn = gr.Button("Help Me Write", size="lg", elem_classes="centered-button")
        def load_grant_details(title):
            if not title:
                return "", "", gr.update(visible=True)
            return "", gr.update(visible=True), gr.update(visible=True)
        def process_grant_details(title):
            if not title:
                return "", "", gr.update(visible=False)
            r = supabase.table('grants').select("*").eq('title', title).execute()
            if not r.data:
                return "", "", gr.update(visible=False)
            g = r.data[0]
            t = f"# {g.get('title', '')}"
            description = g.get('description', '')
            try:
                response = client.chat.completions.create(
                    model="gpt-4",
                    messages=[
                        {"role": "system", "content": "You are a grant writing assistant. Convert the following grant description into clear, concise sentences that highlight the most important information for grant applicants. Focus on key details about the grant's purpose, requirements, and benefits. Keep the language simple and direct."},
                        {"role": "user", "content": description}
                    ],
                    max_tokens=500,
                    temperature=0.3
                )
                processed_description = response.choices[0].message.content
            except:
                processed_description = description
            d = f"""
### Grant Details

**Description**  
{processed_description}

**Important Information**
- Due Date: {g.get('due_date', '')}
- Amount: {g.get('funding_amount', '')}
- Provider: {g.get('provider', '')}
            """
            return t, d, gr.update(visible=False)
        selected_grant.change(
            load_grant_details,
            selected_grant,
            [grant_title, grant_info, loading_indicator]
        ).then(
            process_grant_details,
            selected_grant,
            [grant_title, grant_info, loading_indicator]
        )
        def go_chat():
            return "chat"
        help_write_btn.click(go_chat, None, page_state)

    with gr.Column(visible=False) as chat_page:
        with gr.Row(elem_classes="nav-bar"):
            nav_home_btn3 = gr.Button("Home")
            school_label3 = gr.Markdown("")
        chat_grant_id = gr.Textbox(visible=False)
        chat_school_id = gr.Textbox(visible=False)
        def set_chat_vars(grant_id, school_id):
            print(f"Setting chat variables - Grant ID: {grant_id}, School ID: {school_id}")
            return grant_id, school_id
        def chat_wrapper(message, history):
            grant_id = chat_grant_id.value
            school_id = chat_school_id.value
            print(f"Chat wrapper called with grant_id: {grant_id}, school_id: {school_id}")
            return chat_with_ai(message, history, grant_id, school_id)
        help_write_btn.click(
            set_chat_vars,
            inputs=[selected_grant_id, current_school_id],
            outputs=[chat_grant_id, chat_school_id],
            queue=False
        ).then(
            lambda: "chat",
            None,
            page_state
        )
        chat = gr.ChatInterface(fn=chat_wrapper)

    def go_home():
        return "home"

    nav_home_btn.click(go_home, None, page_state)
    nav_home_btn2.click(go_home, None, page_state)
    nav_home_btn3.click(go_home, None, page_state)

    page_state.change(router, page_state, [landing_page, login_page, signup_page, home_page, grant_details_page, chat_page])

    demo.load(fn=lambda: "landing", inputs=None, outputs=page_state)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True, share=True)