import streamlit as st
import json
import os
import datetime
import pandas as pd
import folium
from streamlit_folium import st_folium
from PIL import Image
from helpers import (
    load_entries, save_entry, get_categories, get_languages,
    text_to_speech, speech_to_text, geocode_location,
    export_to_jsonl, export_to_csv, search_entries,
    register_user, authenticate_user, get_user_info, update_user_entry_count,
    translate_text, detect_language, get_ui_text
)

# Set page config
st.set_page_config(
    page_title="Farming Wisdom Archive",
    page_icon="🌾",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- SESSION STATE INITIALIZATION ---
if 'entries' not in st.session_state:
    st.session_state.entries = load_entries()
if 'authenticated' not in st.session_state:
    st.session_state.authenticated = False
if 'username' not in st.session_state:
    st.session_state.username = None
# NEW: Initialize language selection
if 'language' not in st.session_state:
    st.session_state.language = "English"

# --- LANGUAGE SELECTOR WIDGET ---
# This needs to be defined early so the lang variable is available globally
lang = st.sidebar.selectbox(
    label="Language / भाषा",
    options=["English", "Hindi"],
    key='language'
)

# --- AUTHENTICATION PAGE ---
if not st.session_state.authenticated:
    st.title(get_ui_text("app_title", lang))
    st.markdown(f"#### *{get_ui_text('app_tagline', lang)}*")
    st.markdown("---")
    col1, col2 = st.columns([1, 1])
    with col1:
        st.image("https://images.pexels.com/photos/265216/pexels-photo-265216.jpeg",
                 caption="Photo by Pixabay", use_container_width=True)
        st.markdown(f"### {get_ui_text('about_platform_header', lang)}")
        st.markdown(get_ui_text('about_platform_text', lang))
    with col2:
        tab1, tab2 = st.tabs([f"**{get_ui_text('login_tab', lang)}**", f"**{get_ui_text('register_tab', lang)}**"])
        with tab1:
            st.subheader(get_ui_text('login_subheader', lang))
            with st.form("login_form"):
                username = st.text_input(get_ui_text('username_label', lang))
                password = st.text_input(get_ui_text('password_label', lang), type="password")
                login_button = st.form_submit_button(get_ui_text('login_button', lang), type="primary", use_container_width=True)
                if login_button:
                    if username and password:
                        if authenticate_user(username, password):
                            st.session_state.authenticated = True
                            st.session_state.username = username
                            st.success(get_ui_text('login_success', lang))
                            st.rerun()
                        else:
                            st.error(get_ui_text('invalid_credentials', lang))
                    else:
                        st.error(get_ui_text('enter_credentials', lang))
        with tab2:
            st.subheader(get_ui_text('register_subheader', lang))
            with st.form("register_form"):
                full_name = st.text_input(get_ui_text('full_name_label', lang))
                email = st.text_input(get_ui_text('email_label', lang))
                reg_username = st.text_input(get_ui_text('choose_username_label', lang))
                reg_password = st.text_input(get_ui_text('choose_password_label', lang), type="password")
                confirm_password = st.text_input(get_ui_text('confirm_password_label', lang), type="password")
                register_button = st.form_submit_button(get_ui_text('register_button', lang), type="primary", use_container_width=True)
                if register_button:
                    if full_name and email and reg_username and reg_password and confirm_password:
                        if reg_password == confirm_password:
                            if register_user(reg_username, email, reg_password, full_name):
                                st.success(get_ui_text('registration_success', lang))
                            else:
                                st.error(get_ui_text('registration_failed', lang))
                        else:
                            st.error(get_ui_text('passwords_no_match', lang))
                    else:
                        st.error(get_ui_text('fill_all_fields', lang))
    st.stop()

# --- MAIN APPLICATION (FOR AUTHENTICATED USERS) ---
user_info = get_user_info(st.session_state.username)

# App header with user info
col1, col2 = st.columns([3, 1])
with col1:
    st.title(get_ui_text("app_title", lang))
with col2:
    st.markdown(f"""
    <div style="text-align: right;">
        {get_ui_text('welcome_message', lang)}, <strong>{user_info.get('full_name', st.session_state.username)}</strong>!<br>
        <small>{get_ui_text('entries_submitted', lang)}: {user_info.get('entries_submitted', 0)}</small>
    </div>
    """, unsafe_allow_html=True)
    if st.button(get_ui_text('logout_button', lang), use_container_width=True):
        st.session_state.authenticated = False
        st.session_state.username = None
        st.rerun()

# --- NAVIGATION LOGIC ---
st.sidebar.title(get_ui_text('navigation_header', lang))
def set_page(page_name):
    st.session_state.page_selection = page_name

# Map keys to translated page names
page_keys = {
    "home": get_ui_text("home_page", lang),
    "submit": get_ui_text("submit_page", lang),
    "browse": get_ui_text("browse_page", lang),
    "map": get_ui_text("map_page", lang),
    "search": get_ui_text("search_page", lang),
    "translate": get_ui_text("translate_page", lang),
    "export": get_ui_text("export_page", lang),
    "profile": get_ui_text("profile_page", lang),
}
page_options = list(page_keys.values())

# Get the key of the selected page to use in the if/elif block
selected_page_name = st.sidebar.radio(
    "Go to",
    page_options,
    key="page_selection",
    label_visibility="collapsed"
)
# Find the key corresponding to the selected page name
page_key = [key for key, value in page_keys.items() if value == selected_page_name][0]


# --- Home Page ---
if page_key == "home":
    st.header(get_ui_text("home_header", lang))
    st.markdown("---")
    with st.container(border=True):
        st.subheader(get_ui_text("stats_header", lang))
        total_entries = len(st.session_state.entries)
        languages = len(set(entry.get('language', 'Unknown') for entry in st.session_state.entries))
        categories = len(set(entry.get('category', 'Unknown') for entry in st.session_state.entries))
        
        stat_cols = st.columns(3)
        stat_cols[0].metric(get_ui_text("total_entries_metric", lang), total_entries)
        stat_cols[1].metric(get_ui_text("languages_metric", lang), languages)
        stat_cols[2].metric(get_ui_text("categories_metric", lang), categories)

    st.markdown(f"### {get_ui_text('getting_started_header', lang)}")
    gs_cols = st.columns(3)
    with gs_cols[0]:
        with st.container(border=True, height=180):
            st.markdown(f"#### {get_ui_text('share_knowledge_card_header', lang)}")
            st.markdown(get_ui_text('share_knowledge_card_text', lang))
            st.button(get_ui_text('share_now_button', lang), on_click=set_page, args=(page_keys["submit"],), use_container_width=True)
    with gs_cols[1]:
        with st.container(border=True, height=180):
            st.markdown(f"#### {get_ui_text('explore_entries_card_header', lang)}")
            st.markdown(get_ui_text('explore_entries_card_text', lang))
            st.button(get_ui_text('explore_now_button', lang), on_click=set_page, args=(page_keys["browse"],), use_container_width=True)
    with gs_cols[2]:
        with st.container(border=True, height=180):
            st.markdown(f"#### {get_ui_text('discover_map_card_header', lang)}")
            st.markdown(get_ui_text('discover_map_card_text', lang))
            st.button(get_ui_text('view_map_button', lang), on_click=set_page, args=(page_keys["map"],), use_container_width=True)
            
# --- Submit Wisdom Page ---
elif page_key == "submit":
    st.header(get_ui_text("submit_header", lang))
    st.markdown(get_ui_text("submit_tagline", lang))

    if 'form_data' not in st.session_state:
        st.session_state.form_data = {
            'title': '', 'description': '', 'language': 'English', 'category': 'Other',
            'location_name': '', 'manual_lat': 20.5937, 'manual_lon': 78.9629
        }

    st.subheader(get_ui_text("step1_header", lang))
    with st.container(border=True):
        st.markdown(get_ui_text("pin_location_header", lang))
        st.caption(get_ui_text("pin_location_caption", lang))
        geocode_cols = st.columns([3, 1])
        geocode_input = geocode_cols[0].text_input("search_location", label_visibility="collapsed", placeholder=get_ui_text("find_location_placeholder", lang))
        if geocode_cols[1].button(get_ui_text("find_location_button", lang), use_container_width=True):
            if geocode_input:
                coords = geocode_location(geocode_input)
                if coords:
                    st.session_state.form_data['manual_lat'], st.session_state.form_data['manual_lon'] = coords
                    st.success(get_ui_text("location_found_success", lang).format(location=geocode_input))
                else:
                    st.error(get_ui_text("location_not_found_error", lang).format(location=geocode_input))
        
        st.markdown("---")
        st.markdown(get_ui_text("record_voice_header", lang))
        st.caption(get_ui_text("record_voice_caption", lang))
        speech_cols = st.columns([3, 1])
        speech_language_for_stt = speech_cols[0].selectbox("Language for Speech Recognition", get_languages(), key="speech_lang")
        if speech_cols[1].button(get_ui_text("record_audio_button", lang), use_container_width=True):
             with st.spinner("Listening..."):
                recorded_text = speech_to_text(speech_language_for_stt)
                if recorded_text:
                    st.session_state.form_data['description'] = recorded_text
                    st.success(get_ui_text("speech_recorded_success", lang))
                    st.rerun()
                else:
                    st.error(get_ui_text("speech_record_error", lang))

    st.markdown("---")
    
    st.subheader(get_ui_text("step2_header", lang))
    with st.form("entry_form", clear_on_submit=True):
        with st.expander(get_ui_text("part1_header", lang), expanded=True):
            st.caption(get_ui_text("part1_caption", lang))
            title = st.text_input(get_ui_text("title_label", lang), value=st.session_state.form_data.get('title', ''), placeholder=get_ui_text("title_placeholder", lang))
            description = st.text_area(get_ui_text("description_label", lang), value=st.session_state.form_data.get('description', ''), height=250, placeholder=get_ui_text("description_placeholder", lang))
            
            col1, col2 = st.columns(2)
            category = col1.selectbox(get_ui_text("category_label", lang), get_categories(), index=get_categories().index(st.session_state.form_data.get('category', 'Other')))
            language_of_entry = col2.selectbox(get_ui_text("language_entry_label", lang), get_languages(), index=get_languages().index(st.session_state.form_data.get('language', 'English')))

        with st.expander(get_ui_text("part2_header", lang)):
            st.caption(get_ui_text("part2_caption", lang))
            location_name = st.text_input(get_ui_text("location_name_label", lang), value=st.session_state.form_data.get('location_name', ''))
            map_center = [st.session_state.form_data['manual_lat'], st.session_state.form_data['manual_lon']]
            m = folium.Map(location=map_center, zoom_start=7)
            folium.Marker(map_center, popup="Your Selected Location", tooltip="Current Location").add_to(m)
            map_data = st_folium(m, width=700, height=400)

            if map_data and map_data['last_clicked']:
                st.session_state.form_data['manual_lat'] = map_data['last_clicked']['lat']
                st.session_state.form_data['manual_lon'] = map_data['last_clicked']['lng']
                st.rerun()

            loc_cols = st.columns(2)
            manual_lat = loc_cols[0].text_input(get_ui_text("latitude_label", lang), value=f"{st.session_state.form_data['manual_lat']:.6f}", disabled=True)
            manual_lon = loc_cols[1].text_input(get_ui_text("longitude_label", lang), value=f"{st.session_state.form_data['manual_lon']:.6f}", disabled=True)
            
        with st.expander(get_ui_text("part3_header", lang)):
            st.caption(get_ui_text("part3_caption", lang))
            uploaded_image = st.file_uploader(get_ui_text("image_upload_label", lang), type=['jpg', 'jpeg', 'png'])
            uploaded_audio = st.file_uploader(get_ui_text("audio_upload_label", lang), type=['mp3', 'wav', 'ogg'])

        st.markdown("---")
        submitted = st.form_submit_button(get_ui_text("submit_contribution_button", lang), type="primary", use_container_width=True)

        if submitted:
            if title and description:
                image_path, audio_path = None, None
                media_dir = "data_entries/media"
                os.makedirs(media_dir, exist_ok=True)
                
                if uploaded_image:
                    image_path = os.path.join(media_dir, f"{datetime.datetime.now(datetime.timezone.utc).strftime('%Y%m%d_%H%M%S')}_{uploaded_image.name}")
                    with open(image_path, "wb") as f: f.write(uploaded_image.getbuffer())
                
                if uploaded_audio:
                    audio_path = os.path.join(media_dir, f"{datetime.datetime.now(datetime.timezone.utc).strftime('%Y%m%d_%H%M%S')}_{uploaded_audio.name}")
                    with open(audio_path, "wb") as f: f.write(uploaded_audio.getbuffer())

                entry = {
                    'id': len(st.session_state.entries) + 1, 'title': title, 'description': description,
                    'language': language_of_entry, 'category': category, 'location_name': location_name,
                    'latitude': st.session_state.form_data['manual_lat'],
                    'longitude': st.session_state.form_data['manual_lon'],
                    'image_path': image_path, 'audio_path': audio_path,
                    'timestamp': datetime.datetime.now(datetime.timezone.utc).isoformat(),
                    'contributor': st.session_state.username,
                    'contributor_full_name': user_info.get('full_name', st.session_state.username)
                }
                
                if save_entry(entry):
                    st.session_state.entries.append(entry)
                    update_user_entry_count(st.session_state.username)
                    st.success(get_ui_text("submit_success", lang))
                    st.session_state.pop('form_data', None)
                else:
                    st.error(get_ui_text("submit_failed", lang))
            else:
                st.error(get_ui_text("submit_validation_error", lang))

# --- Browse Knowledge Page ---
elif page_key == "browse":
    st.header(get_ui_text("browse_header", lang))
    st.markdown(get_ui_text("browse_tagline", lang))
    
    with st.container(border=True):
        col1, col2, col3 = st.columns(3)
        with col1:
            filter_language = st.selectbox(get_ui_text("filter_lang_label", lang), [get_ui_text("all_option", lang)] + get_languages())
        with col2:
            filter_category = st.selectbox(get_ui_text("filter_cat_label", lang), [get_ui_text("all_option", lang)] + get_categories())
        with col3:
            sort_by = st.selectbox(get_ui_text("sort_by_label", lang), [get_ui_text("newest_first_option", lang), get_ui_text("oldest_first_option", lang), get_ui_text("title_az_option", lang)])

    filtered_entries = st.session_state.entries.copy()
    if filter_language != get_ui_text("all_option", lang): filtered_entries = [e for e in filtered_entries if e.get('language') == filter_language]
    if filter_category != get_ui_text("all_option", lang): filtered_entries = [e for e in filtered_entries if e.get('category') == filter_category]
    
    if sort_by == get_ui_text("newest_first_option", lang): filtered_entries.sort(key=lambda x: x.get('timestamp', ''), reverse=True)
    elif sort_by == get_ui_text("oldest_first_option", lang): filtered_entries.sort(key=lambda x: x.get('timestamp', ''))
    elif sort_by == get_ui_text("title_az_option", lang): filtered_entries.sort(key=lambda x: x.get('title', '').lower())

    st.markdown(get_ui_text("showing_entries", lang).format(count=len(filtered_entries), total=len(st.session_state.entries)))
    st.markdown("---")

    for entry in filtered_entries:
        with st.container(border=True):
            st.subheader(f"📖 {entry.get('title', 'Untitled')}")
            
            meta_cols = st.columns(4)
            meta_cols[0].markdown(get_ui_text("category_display", lang).format(category=entry.get('category', 'N/A')))
            meta_cols[1].markdown(get_ui_text("language_display", lang).format(language=entry.get('language', 'N/A')))
            meta_cols[2].markdown(get_ui_text("location_display", lang).format(location=entry.get('location_name', 'N/A')))
            meta_cols[3].markdown(get_ui_text("submitted_display", lang).format(date=entry.get('timestamp', 'N/A')[:10]))

            with st.expander(get_ui_text("view_details_expander", lang)):
                col1, col2 = st.columns([2, 1])
                with col1:
                    st.markdown(get_ui_text("description_display", lang))
                    st.write(entry.get('description', 'No description'))
                    if st.button(get_ui_text("listen_button", lang), key=f"tts_{entry.get('id')}"):
                        text_to_speech(entry.get('description', ''), entry.get('language'))
                
                with col2:
                    if entry.get('image_path') and os.path.exists(entry['image_path']):
                        st.image(Image.open(entry['image_path']), caption=get_ui_text("image_caption", lang), use_container_width=True)
                    if entry.get('audio_path') and os.path.exists(entry['audio_path']):
                        st.audio(entry['audio_path'])
                    if entry.get('latitude') and entry.get('longitude'):
                        st.markdown(get_ui_text("coords_display", lang).format(lat=f"{entry['latitude']:.4f}", lon=f"{entry['longitude']:.4f}"))

# --- Knowledge Map Page ---
elif page_key == "map":
    st.header(get_ui_text("map_header", lang))
    st.markdown(get_ui_text("map_tagline", lang))
    
    geo_entries = [e for e in st.session_state.entries if e.get('latitude') and e.get('longitude')]
    
    if geo_entries:
        m = folium.Map(location=[20.5937, 78.9629], zoom_start=5)
        for entry in geo_entries:
            popup_html = f"""
            <b>{entry.get('title', 'Untitled')}</b><br>
            <b>{get_ui_text("map_popup_category", lang)}</b> {entry.get('category', 'N/A')}<br>
            <b>{get_ui_text("map_popup_language", lang)}</b> {entry.get('language', 'N/A')}
            """
            folium.Marker(
                [entry['latitude'], entry['longitude']],
                popup=popup_html,
                tooltip=entry.get('title', 'Untitled')
            ).add_to(m)
        st_folium(m, width=1200, height=600, returned_objects=[])
        st.info(get_ui_text("map_info", lang).format(count=len(geo_entries)))
    else:
        st.warning(get_ui_text("map_warning", lang))

# --- Search Page ---
elif page_key == "search":
    st.header(get_ui_text("search_header", lang))
    st.markdown(get_ui_text("search_tagline", lang))

    with st.container(border=True):
        search_query = st.text_input("search", label_visibility="collapsed", placeholder=get_ui_text("search_placeholder", lang))
        with st.expander(get_ui_text("advanced_filters_expander", lang)):
            col1, col2 = st.columns(2)
            search_language = col1.selectbox(get_ui_text("filter_lang_label", lang), [get_ui_text("all_option", lang)] + get_languages(), key="search_lang")
            search_category = col2.selectbox(get_ui_text("filter_cat_label", lang), [get_ui_text("all_option", lang)] + get_categories(), key="search_cat")
            has_media = col1.checkbox(get_ui_text("has_media_checkbox", lang))
            has_location = col2.checkbox(get_ui_text("has_location_checkbox", lang))

    if search_query:
        results = search_entries(
            st.session_state.entries, search_query,
            language=search_language,
            category=search_category,
            has_media=has_media, has_location=has_location
        )
        st.markdown(get_ui_text("search_results_header", lang).format(count=len(results), query=search_query))
        for entry in results:
            with st.container(border=True):
                 st.subheader(f"📖 {entry.get('title', 'Untitled')}")
                 st.markdown(f"**{get_ui_text('map_popup_category', lang)}** {entry.get('category', 'N/A')} | **{get_ui_text('map_popup_language', lang)}** {entry.get('language', 'N/A')} | **{get_ui_text('location_display', lang).format(location=entry.get('location_name', 'N/A'))}**")
                 st.write(entry.get('description', 'No description'))

# --- Translation Hub Page ---
elif page_key == "translate":
    st.header(get_ui_text("translation_hub_header", lang))
    st.markdown(get_ui_text("translation_hub_tagline", lang))

    with st.container(border=True):
        col1, col2 = st.columns(2)
        with col1:
            st.subheader(get_ui_text("original_text_header", lang))
            source_language = st.selectbox(get_ui_text("source_language_label", lang), ["Auto-detect"] + get_languages())
            source_text = st.text_area(get_ui_text("text_to_translate_label", lang), height=200, placeholder=get_ui_text("text_to_translate_placeholder", lang), label_visibility="collapsed")

        with col2:
            st.subheader(get_ui_text("translated_text_header", lang))
            target_language = st.selectbox(get_ui_text("target_language_label", lang), get_languages())
            if st.button(get_ui_text("translate_button", lang), type="primary", use_container_width=True) and source_text:
                with st.spinner("Translating..."):
                    source_lang_code = source_language if source_language != "Auto-detect" else "auto"
                    translated_text = translate_text(source_text, target_language, source_lang_code)
                    st.text_area(get_ui_text("translation_output_label", lang), value=translated_text, height=200, key="translated_output", label_visibility="collapsed")
            else:
                 st.text_area(get_ui_text("translation_output_label", lang), value="", height=200, label_visibility="collapsed")

# --- Export Data Page ---
elif page_key == "export":
    st.header(get_ui_text("export_header", lang))
    st.markdown(get_ui_text("export_tagline", lang))

    with st.container(border=True):
        st.subheader(get_ui_text("export_config_header", lang))
        col1, col2 = st.columns(2)
        export_format = col1.selectbox(get_ui_text("format_label", lang), ["CSV", "JSONL"])
        export_language = col2.selectbox(get_ui_text("filter_lang_label", lang), [get_ui_text("all_option", lang)] + get_languages(), key="export_lang")
        export_category = col1.selectbox(get_ui_text("filter_cat_label", lang), [get_ui_text("all_option", lang)] + get_categories(), key="export_cat")
        
        if st.button(get_ui_text("generate_export_button", lang), type="primary", use_container_width=True):
            filtered_entries = st.session_state.entries.copy()
            if export_language != get_ui_text("all_option", lang): filtered_entries = [e for e in filtered_entries if e.get('language') == export_language]
            if export_category != get_ui_text("all_option", lang): filtered_entries = [e for e in filtered_entries if e.get('category') == export_category]

            if not filtered_entries:
                st.warning(get_ui_text("export_no_match_warning", lang))
            else:
                timestamp = datetime.datetime.now(datetime.timezone.utc).strftime('%Y%m%d')
                if export_format == "CSV":
                    export_data = export_to_csv(filtered_entries, True, True)
                    st.download_button(
                        label=get_ui_text("export_download_csv_button", lang), data=export_data,
                        file_name=f"farming_wisdom_{timestamp}.csv", mime="text/csv", use_container_width=True
                    )
                else: # JSONL
                    export_data = export_to_jsonl(filtered_entries, True, True)
                    st.download_button(
                        label=get_ui_text("export_download_jsonl_button", lang), data=export_data,
                        file_name=f"farming_wisdom_{timestamp}.jsonl", mime="application/jsonl", use_container_width=True
                    )
                st.success(get_ui_text("export_ready_success", lang).format(count=len(filtered_entries)))

# --- Profile Page ---
elif page_key == "profile":
    st.header(get_ui_text("profile_header", lang).format(name=user_info.get('full_name', st.session_state.username)))
    st.markdown("---")
    
    tab1, tab2, tab3 = st.tabs([f"**{get_ui_text('my_info_tab', lang)}**", f"**{get_ui_text('my_contributions_tab', lang)}**", f"**{get_ui_text('settings_tab', lang)}**"])

    with tab1:
        with st.container(border=True):
            st.subheader(get_ui_text("profile_info_header", lang))
            info_cols = st.columns(2)
            info_cols[0].markdown(get_ui_text("full_name_display", lang).format(name=user_info.get('full_name', 'Not provided')))
            info_cols[0].markdown(get_ui_text("username_display", lang).format(username=st.session_state.username))
            info_cols[1].markdown(get_ui_text("email_display", lang).format(email=user_info.get('email', 'Not provided')))
            info_cols[1].markdown(get_ui_text("member_since_display", lang).format(date=user_info.get('registration_date', 'Unknown')[:10]))
            st.metric(get_ui_text("total_entries_metric_profile", lang), user_info.get('entries_submitted', 0))

    with tab2:
        with st.container(border=True):
            st.subheader(get_ui_text("my_recent_contributions_header", lang))
            my_entries = [e for e in st.session_state.entries if e.get('contributor') == st.session_state.username]
            if my_entries:
                st.write(get_ui_text("contributions_count_message", lang).format(count=len(my_entries)))
                for entry in reversed(my_entries[-5:]):
                    st.markdown(f"- **{entry.get('title', 'Untitled')}** (*{entry.get('category', 'N/A')}*)")
            else:
                st.info(get_ui_text("no_contributions_message", lang))

    with tab3:
        with st.container(border=True):
            st.subheader(get_ui_text("account_prefs_header", lang))
            st.selectbox(get_ui_text("preferred_language_label", lang), get_languages(), index=1)
            st.checkbox(get_ui_text("email_notifications_checkbox", lang), value=False)
            if st.button(get_ui_text("update_profile_button", lang), use_container_width=True):
                st.success(get_ui_text("profile_update_success", lang))

# --- Footer ---
st.sidebar.markdown("---")
st.sidebar.info(get_ui_text("app_version_footer", lang))