import os
import streamlit as st
import json
from streamlit_option_menu import option_menu
from gemini_utility import (load_gemini_pro, gemini_pro_vision_responce)
from PIL import Image
import time

# Configuración de la página
st.set_page_config(
    page_title="GnosticDev AI",
    page_icon="🤖",
    layout="centered",
    initial_sidebar_state="expanded",
)

# Archivo para almacenamiento permanente
PROMPTS_FILE = "training_prompts.json"

# Funciones para manejar el almacenamiento permanente
def load_prompts():
    try:
        if os.path.exists(PROMPTS_FILE):
            with open(PROMPTS_FILE, 'r', encoding='utf-8') as file:
                return json.load(file)
    except Exception as e:
        st.error(f"Error cargando prompts: {e}")
    return {"prompts": [], "current_prompt": ""}

def save_prompts(prompts_data):
    try:
        with open(PROMPTS_FILE, 'w', encoding='utf-8') as file:
            json.dump(prompts_data, file, ensure_ascii=False, indent=2)
    except Exception as e:
        st.error(f"Error guardando prompts: {e}")

# Inicialización de datos
if 'prompts_data' not in st.session_state:
    st.session_state.prompts_data = load_prompts()

with st.sidebar:
    selected = option_menu(
        "GD AI",
        ["System Prompts", "Chatbot", "Image Captioning"],
        menu_icon="robot",
        icons=['gear', 'chat-dots-fill', 'image-fill'],
        default_index=0
    )

if selected == "System Prompts":
    st.title("Gestión de System Prompts")
    
    # Área para nuevo prompt
    new_prompt = st.text_area(
        "Nuevo System Prompt",
        value="",
        height=200,
        help="Escribe aquí las nuevas instrucciones para el AI"
    )
    
    col1, col2 = st.columns([1, 2])
    with col1:
        if st.button("Añadir Nuevo Prompt"):
            if new_prompt and new_prompt not in st.session_state.prompts_data["prompts"]:
                st.session_state.prompts_data["prompts"].append(new_prompt)
                save_prompts(st.session_state.prompts_data)
                st.success("Nuevo prompt añadido!")
                time.sleep(1)
                st.rerun()
    
    # Lista de prompts guardados
    st.markdown("### Prompts Guardados")
    for i, prompt in enumerate(st.session_state.prompts_data["prompts"]):
        with st.expander(f"Prompt {i+1}"):
            st.text_area("", prompt, height=100, key=f"prompt_{i}", disabled=True)
            col1, col2, col3 = st.columns([1, 1, 1])
            with col1:
                if st.button("Usar este prompt", key=f"use_{i}"):
                    st.session_state.prompts_data["current_prompt"] = prompt
                    save_prompts(st.session_state.prompts_data)
                    if "chat_session" in st.session_state:
                        del st.session_state.chat_session
                    st.success("Prompt activado!")
            with col2:
                if st.button("Editar", key=f"edit_{i}"):
                    st.session_state.editing_prompt = i
                    st.session_state.editing_text = prompt
            with col3:
                if st.button("Eliminar", key=f"delete_{i}"):
                    st.session_state.prompts_data["prompts"].pop(i)
                    save_prompts(st.session_state.prompts_data)
                    st.success("Prompt eliminado!")
                    time.sleep(1)
                    st.rerun()
    
    # Mostrar prompt actual
    st.markdown("### Prompt Actual")
    current_prompt = st.session_state.prompts_data.get("current_prompt", "")
    if current_prompt:
        st.info(current_prompt)
    else:
        st.warning("No hay prompt activo")

elif selected == "Chatbot":
    model = load_gemini_pro()
    
    # Initialize chat session with current prompt
    if "chat_session" not in st.session_state:
        st.session_state.chat_session = model.start_chat(history=[])
        current_prompt = st.session_state.prompts_data.get("current_prompt")
        if current_prompt:
            st.session_state.chat_session.send_message(current_prompt)

    st.title("Gnosticdev Chatbot")
    
    # Mostrar prompt actual
    current_prompt = st.session_state.prompts_data.get("current_prompt")
    if current_prompt:
        with st.expander("Ver System Prompt actual"):
            st.info(current_prompt)
    
    # Display chat history
    for message in st.session_state.chat_session.history:
        with st.chat_message(translate_role_to_streamlit(message.role)):
            st.markdown(message.parts[0].text)

    # Chat input
    user_prompt = st.chat_input("Preguntame algo...")
    if user_prompt:
        st.chat_message("user").markdown(user_prompt)
        gemini_response = st.session_state.chat_session.send_message(user_prompt)
        with st.chat_message("assistant"):
            st.markdown(gemini_response.text)

elif selected == "Image Captioning":
    st.title("Image Caption Generation📸")
    upload_image = st.file_uploader("Upload an image...", type=["jpg", "jpeg", "png"])
    if upload_image and st.button("Generate"):
        image = Image.open(upload_image)
        col1, col2 = st.columns(2)
        with col1:
            st.image(image, caption="Uploaded Image", use_column_width=True)
        default_prompt = "Write a caption for this image"
        caption = gemini_pro_vision_responce(default_prompt, image)
        with col2:
            st.info(caption)