from pathlib import Path
from PIL import Image
import streamlit as st
import insightface
from insightface.app import FaceAnalysis
from huggingface_hub import InferenceClient, AsyncInferenceClient
import asyncio
import os
import random
import numpy as np
import yaml

try:
    with open("config.yaml", "r") as file:
        credentials = yaml.safe_load(file)
except Exception as e:
    st.error(f"Error al cargar el archivo de configuración: {e}")
    credentials = {"username": "", "password": ""}

MAX_SEED = np.iinfo(np.int32).max
client = AsyncInferenceClient()
llm_client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
DATA_PATH = Path("./data")
DATA_PATH.mkdir(exist_ok=True)
PREDEFINED_SEED = random.randint(0, MAX_SEED)
HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")

if not HF_TOKEN_UPSCALER:
    st.warning("HF_TOKEN_UPSCALER no está configurado. Algunas funcionalidades pueden no funcionar.")

def login_form():
    st.title("Iniciar Sesión")
    username = st.text_input("Usuario", value="admin")
    password = st.text_input("Contraseña", value="flux3x", type="password") 
    if st.button("Iniciar Sesión"):
        if authenticate_user(username, password):
            st.success("Autenticación exitosa.")
            st.session_state['authenticated'] = True
        else:
            st.error("Credenciales incorrectas. Intenta de nuevo.")

def get_upscale_finegrain(prompt, img_path, upscale_factor):
    try:
        upscale_client = InferenceClient("fal/AuraSR-v2", hf_token=HF_TOKEN_UPSCALER)
        result = upscale_client.predict(input_image=handle_file(img_path), prompt=prompt, upscale_factor=upscale_factor)
        return result[1] if isinstance(result, list) and len(result) > 1 else None
    except Exception as e:
        st.error(f"Error al mejorar la imagen: {e}")
        return None

def authenticate_user(username, password):
    return username == credentials["username"] and password == credentials["password"]

def prepare_face_app():
    app = FaceAnalysis(name='buffalo_l')
    app.prepare(ctx_id=0, det_size=(640, 640))
    swapper = insightface.model_zoo.get_model('onix.onnx')
    return app, swapper

app, swapper = prepare_face_app()

def sort_faces(faces):
    return sorted(faces, key=lambda x: x.bbox[0])

def get_face(faces, face_id):
    if not faces or len(faces) < face_id:
        raise ValueError("Rostro no disponible.")
    return faces[face_id - 1]

def swap_faces(source_image, source_face_index, destination_image, destination_face_index):
    faces = sort_faces(app.get(source_image))
    source_face = get_face(faces, source_face_index)
    
    res_faces = sort_faces(app.get(destination_image))
    if destination_face_index > len(res_faces) or destination_face_index < 1:
        raise ValueError("Índice de rostro de destino no válido.")
    
    res_face = get_face(res_faces, destination_face_index)
    result = swapper.get(destination_image, res_face, source_face, paste_back=True)
    return result

def generate_image(prompt, width, height, seed, model_name):
    if seed == -1:
        seed = random.randint(0, MAX_SEED)
    image = client.text_to_image(prompt=prompt, height=height, width=width, model=model_name)
    return image, seed

def gen(prompts, width, height, model_name, num_variants=1, use_enhanced=True):
    images = []
    try:
        for idx, prompt in enumerate(prompts):
            seed = random.randint(0, MAX_SEED)
            image, seed = generate_image(prompt, width, height, seed, model_name)
            image_path = save_image(image, f"generated_image_{seed}.jpg")
            if image_path:
                st.success(f"Imagen {idx + 1} generada")
                images.append(str(image_path))
    except Exception as e:
        st.error(f"Error al generar imágenes: {e}")
    return images

def improve_prompt(prompt):
    try:
        instructions = [
            "With my idea create a vibrant description for a detailed txt2img prompt, 300 characters max.",
            "With my idea write a creative and detailed text-to-image prompt in English, 300 characters max.",
            "With my idea generate a descriptive and visual txt2img prompt in English, 300 characters max.",
            "With my idea describe a photorealistic with illumination txt2img prompt in English, 300 characters max.",
            "With my idea give a realistic and elegant txt2img prompt in English, 300 characters max.",
            "With my idea conform a visually dynamic and surreal txt2img prompt in English, 300 characters max.",
            "With my idea realize an artistic and cinematic txt2img prompt in English, 300 characters max."
        ]
        instruction = random.choice(instructions)
        formatted_prompt = f"{prompt}: {instruction}"
        response = llm_client.text_generation(formatted_prompt, max_new_tokens=100)
        return response['generated_text'][:100] if 'generated_text' in response else response.strip()
    except Exception as e:
        return f"Error mejorando el prompt: {e}"

def generate_variations(prompt, num_variants, use_enhanced):
    prompts = set() 
    while len(prompts) < num_variants:
        if use_enhanced:
            enhanced_prompt = improve_prompt(prompt)
            prompts.add(enhanced_prompt)
        else:
            prompts.add(prompt)
    return list(prompts)

def delete_all_images():
    try:
        for image_file in list_saved_images():
            os.remove(image_file)
        st.success("Todas las imágenes han sido eliminadas.")
    except Exception as e:
        st.error(f"Error al eliminar las imágenes: {e}")

def upload_images_to_gallery():
    uploaded_images = st.sidebar.file_uploader("Sube imágenes a la galería", type=["jpg", "jpeg", "png"], accept_multiple_files=True)
    if uploaded_images:
        for uploaded_image in uploaded_images:
            image = Image.open(uploaded_image)
            image_path = save_image(image, f"{uploaded_image.name}")
            if image_path:
                save_prompt("uploaded by user")
                st.sidebar.success(f"Imagen subida: {image_path}")

def main():
    st.set_page_config(layout="wide")

    if 'authenticated' not in st.session_state or not st.session_state['authenticated']:
        login_form()
        return

    st.title("Flux +Upscale +Prompt Enhancer +FaceSwap")
    
    if st.sidebar.button("Borrar todas las imágenes"):
        delete_all_images()

    generated_image_path = st.session_state.get('generated_image_path')
    prompt = st.sidebar.text_area("Descripción de la imagen", height=150, max_chars=500)
    format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
    model_option = st.sidebar.selectbox("Modelo", ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev", "enhanceaiteam/Flux-Uncensored-V2", "enhanceaiteam/Flux-Uncensored"])
    prompt_checkbox = st.sidebar.checkbox("Prompt Enhancer")
    upscale_checkbox = st.sidebar.checkbox("Escalar imagen")
    
    width, height = (360, 640) if format_option == "9:16" else (640, 360) if format_option == "16:9" else (640, 640)
    num_variants = st.sidebar.slider("Número de imágenes a generar", 1, 8, 1) if prompt_checkbox else 1
    
    if prompt_checkbox:
        with st.spinner("Generando prompts mejorados..."):
            prompts = generate_variations(prompt, num_variants, True)
    else:
        prompts = [prompt] 

    upload_images_to_gallery()

    if st.sidebar.button("Generar Imágenes"):
        with st.spinner("Generando imágenes..."):
            try:
                results = gen(prompts, width, height, model_option, num_variants, prompt_checkbox)
                st.session_state['generated_image_paths'] = results
                for result in results:
                    st.image(result, caption="Imagen Generada")
            except Exception as e:
                st.error(f"Error al generar las imágenes: {str(e)}")

    if generated_image_path:
        if upscale_checkbox:
            with st.spinner("Escalando imagen..."):
                try:
                    upscale_image_path = get_upscale_finegrain("Upscale", generated_image_path, 2)
                    if upscale_image_path:
                        st.image(upscale_image_path, caption="Imagen Escalada")
                except Exception as e:
                    st.error(f"Error al escalar la imagen: {str(e)}")

        st.header("Intercambio de Rostros")
        source_image_file = st.file_uploader("Imagen de Origen", type=["jpg", "jpeg", "png"])

        if source_image_file is not None:
            try:
                source_image = Image.open(source_image_file)
            except Exception as e:
                st.error(f"Error al cargar la imagen de origen: {str(e)}")
                source_image = None
        else:
            source_image = Image.open("face.jpg")

        source_face_index = st.number_input('Posición del Rostro', min_value=1, value=1, key="source_face_index")
        destination_face_index = st.number_input('Posición del Rostro de Destino', min_value=1, value=1, key="destination_face_index")

        if st.button("Intercambiar Rostros"):
            try:
                destination_image = Image.open(generated_image_path)
                result_image = swap_faces(np.array(source_image), source_face_index, np.array(destination_image), destination_face_index)
                swapped_image = Image.fromarray(result_image)
                
                swapped_image_path = save_image(swapped_image, f"swapped_image_{PREDEFINED_SEED}.jpg")
                if swapped_image_path:
                    st.image(swapped_image, caption="Intercambio de Rostro")
                    os.remove(generated_image_path)
                else:
                    st.warning("La imagen intercambiada ya existe en la galería.")
            except Exception as e:
                st.error(f"Ocurrió un error al intercambiar rostros: {str(e)}")

    display_gallery()

if __name__ == "__main__":
    main()