import streamlit as st
import google.generativeai as genai
import os
import time
from dotenv import load_dotenv
from styles import get_custom_css, get_response_html_wrapper
from formulas import offer_formulas
import PyPDF2
import docx
from PIL import Image
import io

# Import the bullet generator
from bullets.generator import create_bullet_instruction

# Set page to wide mode to use full width
st.set_page_config(layout="wide")

# Load environment variables
load_dotenv()

# Configure Google Gemini API
genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
model = genai.GenerativeModel('gemini-2.0-flash')

# Import the create_offer_instruction function from formulas
from formulas import create_offer_instruction, offer_formulas

# Initialize session state variables if they don't exist
if 'submitted' not in st.session_state:
    st.session_state.submitted = False
if 'offer_result' not in st.session_state:
    st.session_state.offer_result = ""
if 'generated' not in st.session_state:
    st.session_state.generated = False

# Hide Streamlit menu and footer
st.markdown("""
<style>
#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
header {visibility: hidden;}
</style>
""", unsafe_allow_html=True)

# Custom CSS
st.markdown(get_custom_css(), unsafe_allow_html=True)

# App title and description
st.markdown('<h1 style="text-align: center;">Great Offer Generator</h1>', unsafe_allow_html=True)
st.markdown('<h3 style="text-align: center;">Transform your skills into compelling offers!</h3>', unsafe_allow_html=True)

# Create two columns for layout - left column 40%, right column 60%
col1, col2 = st.columns([4, 6])

# Main input section in left column
with col1:
    # Define the generate_offer function first
    def handle_generate_button():  # Renamed to avoid conflict
        has_manual_input = bool(skills or product_service)
        has_file_input = bool(uploaded_file is not None and not is_image)
        has_image_input = bool(uploaded_file is not None and is_image)

        # Simple validation - check if we have at least one input type
        if not (has_manual_input or has_file_input or has_image_input):
            st.error('Por favor ingresa texto o sube un archivo/imagen')
            return
        
        st.session_state.submitted = True
        st.session_state.generated = False  # Reset generated flag
        
        # Store inputs based on what's available
        if has_manual_input:
            st.session_state.skills = skills if skills else ""
            st.session_state.product_service = product_service if product_service else ""
        
        if has_file_input:
            st.session_state.file_content = file_content
            
        if has_image_input:
            st.session_state.image_parts = image_parts
        
        # Set input type based on what's available
        if has_image_input:
            if has_manual_input:
                st.session_state.input_type = "manual_image"
            else:
                st.session_state.input_type = "image"
        else:
            if has_manual_input and has_file_input:
                st.session_state.input_type = "both"
            elif has_file_input:
                st.session_state.input_type = "file"
            elif has_manual_input:
                st.session_state.input_type = "manual"
        
        # Store common settings
        st.session_state.target_audience = target_audience
        st.session_state.temperature = temperature
        st.session_state.formula_type = formula_type

    # Keep only the manual input tab
    with st.container():
        skills = st.text_area('💪 Tus Habilidades', height=70, 
                            help='Lista tus habilidades y experiencia clave')
        product_service = st.text_area('🎯 Producto/Servicio', height=70,
                                    help='Describe tu producto o servicio')
        
        # Generate button moved here - right after product/service
        st.button('Generar Oferta 🎉', on_click=handle_generate_button)  # Updated function name
    
    # Accordion for additional settings
    with st.expander('⚙️ Configuración Avanzada'):
        target_audience = st.text_area('👥 Público Objetivo', height=70,
                                    help='Describe tu cliente o público ideal')
        
        # Add file/image uploader here
        uploaded_file = st.file_uploader("📄 Sube un archivo o imagen", 
                                        type=['txt', 'pdf', 'docx', 'jpg', 'jpeg', 'png'])
        
        if uploaded_file is not None:
            file_type = uploaded_file.name.split('.')[-1].lower()
            
            # Handle text files
            if file_type in ['txt', 'pdf', 'docx']:
                if file_type == 'txt':
                    try:
                        file_content = uploaded_file.read().decode('utf-8')
                    except Exception as e:
                        st.error(f"Error al leer el archivo TXT: {str(e)}")
                        file_content = ""
                    
                elif file_type == 'pdf':
                    try:
                        import PyPDF2
                        pdf_reader = PyPDF2.PdfReader(uploaded_file)
                        file_content = ""
                        for page in pdf_reader.pages:
                            file_content += page.extract_text() + "\n"
                    except Exception as e:
                        st.error(f"Error al leer el archivo PDF: {str(e)}")
                        file_content = ""
                        
                elif file_type == 'docx':
                    try:
                        import docx
                        doc = docx.Document(uploaded_file)
                        file_content = "\n".join([para.text for para in doc.paragraphs])
                    except Exception as e:
                        st.error(f"Error al leer el archivo DOCX: {str(e)}")
                        file_content = ""
                
                # Remove success message - no notification shown
                
                # Set file type flag
                is_image = False
            
            # Handle image files
            elif file_type in ['jpg', 'jpeg', 'png']:
                try:
                    image = Image.open(uploaded_file)
                    st.image(image, caption="Imagen cargada", use_container_width=True)
                    
                    image_bytes = uploaded_file.getvalue()
                    image_parts = [
                        {
                            "mime_type": uploaded_file.type,
                            "data": image_bytes
                        }
                    ]
                    
                    # Set file type flag
                    is_image = True
                except Exception as e:
                    st.error(f"Error al procesar la imagen: {str(e)}")
                    is_image = False
        
        # Selector de fórmula
        formula_type = st.selectbox(
            '📋 Tipo de Fórmula',
            options=list(offer_formulas.keys()),
            help='Selecciona el tipo de fórmula para tu oferta'
        )
        
        temperature = st.slider('🌡️ Nivel de Creatividad', min_value=0.0, max_value=2.0, value=1.0,
                            help='Valores más altos hacen que el resultado sea más creativo pero menos enfocado')

# Results column
# In the section where you're generating the offer
with col2:
    if st.session_state.submitted and not st.session_state.generated:
        with st.spinner('Creando tu oferta perfecta...'):
            # Use the create_offer_instruction function to generate the prompt
            avatar_description = st.session_state.target_audience if hasattr(st.session_state, 'target_audience') and st.session_state.target_audience else 'General audience'
            
            # Determine product name based on input type
            if hasattr(st.session_state, 'product_service') and st.session_state.product_service:
                product_name = st.session_state.product_service
            else:
                product_name = "Producto/Servicio"
                
            # Get the instruction using the formula
            instruction = create_offer_instruction(
                avatar_description=avatar_description,
                product_name=product_name,
                selected_formula_name=st.session_state.formula_type
            )
            
            # Add specific instructions for more creative hooks when using Oferta Dorada
            if st.session_state.formula_type == "Oferta Dorada":
                creative_hook_instruction = """
                IMPORTANT INSTRUCTION FOR THE ATTENTION HOOK:
                Make the first line (attention hook) extremely creative, disruptive, and emotionally impactful.
                Use unexpected comparisons, shocking statistics, provocative questions, or dramatic scenarios.
                The hook should be surprising and immediately grab attention with unconventional language or framing.
                Avoid generic statements and create a sense of urgency or curiosity.
                """
                instruction += creative_hook_instruction
            
            # Add instruction for generating benefit bullets based on the promise
            instruction += create_bullet_instruction()
            
            # Add instruction for generating bonuses that complement the offer
            from formulas import create_bonus_instruction
            bonus_instruction = create_bonus_instruction(
                avatar_description=avatar_description,
                product_name=product_name,
                selected_formula_name=st.session_state.formula_type
            )
            instruction += "\n\n" + bonus_instruction
            
            # Add additional context based on input type
            if st.session_state.input_type == "manual":
                additional_context = f"""
                Additional information:
                Skills: {st.session_state.skills}
                """
                instruction += additional_context
            elif st.session_state.input_type == "file":
                additional_context = f"""
                Additional information from file:
                {st.session_state.file_content}
                """
                instruction += additional_context
            elif st.session_state.input_type == "both":
                additional_context = f"""
                Additional information:
                Skills: {st.session_state.skills}
                File content: {st.session_state.file_content}
                """
                instruction += additional_context
                
            try:
                generation_config = genai.GenerationConfig(temperature=st.session_state.temperature)
                
                if "image" in st.session_state.input_type:
                    response = model.generate_content([instruction, st.session_state.image_parts[0]], generation_config=generation_config)
                else:
                    response = model.generate_content(instruction, generation_config=generation_config)
                    
                st.session_state.offer_result = response.text
                st.session_state.generated = True  # Mark as generated
                
            except Exception as e:
                st.error(f'Ocurrió un error: {str(e)}')
                st.session_state.submitted = False
    
    # Display results if we have an offer result
    if st.session_state.generated:
        # With this line that uses the wrapper:
        st.markdown(get_response_html_wrapper(st.session_state.offer_result), unsafe_allow_html=True)
        
        # Add a small space
        st.markdown('<div style="height: 15px;"></div>', unsafe_allow_html=True)
        
        # Apply the custom button style before rendering the download button
        st.markdown('<style>div.stDownloadButton > button {your-custom-styles-here}</style>', unsafe_allow_html=True)
        st.download_button(
            label="Descargar Oferta",
            data=st.session_state.offer_result,
            file_name="oferta_generada.txt",
            mime="text/plain"
        )

# Footer
st.markdown('---')
st.markdown('Made with ❤️ by Jesús Cabrera')

# Remove the duplicate functions at the bottom