import os
import requests
import json
import html2text
import tqdm
import re
import gradio as gr
import modules.scripts as scripts
from modules import script_callbacks
from extensions.civitai.scripts.updater import *
import hashlib
import pandas as pd

class Script(scripts.Script):
    def __init__(self) -> None:
        super().__init__()
        self.html2text = html2text.HTML2Text()

    def title(self):
        return "Civitai"

    def show(self, is_img2img):
        return scripts.AlwaysVisible

    def ui(self, is_img2img):
        return ()

    def on_ui_tabs(self):
        return [(self.interface, "Civitai", "civitai")]

directories = {
    "SD Checkpoints": "models/Stable-diffusion/",
    "SD Loras": "models/Lora/",
    "Flux Checkpoints": "models/Stable-diffusion/Flux",
    "Flux Loras": "models/Lora/Flux"
}

api = "1f268e312e5b8af60ff0e86acd959236"

class CivitaiDownloader:
    def __init__(self):
        self.html2text = html2text.HTML2Text()
        self.selected_image_url = None
        self.stop_download = False  # Variable de control para detener la descarga

    def extract_model_id(self, url):
        try:
            parts = url.strip().split('/')
            model_index = parts.index('models') + 1
            return parts[model_index]
        except (IndexError, ValueError):
            return None

    def get_model_info(self, url, api_token):
        model_id = self.extract_model_id(url)
        if not model_id:
            return None, None, "URL de Civitai inválida."

        api_url = f'https://civitai.com/api/v1/models/{model_id}'
        headers = {'Authorization': f'Bearer {api_token}'} if api_token else {}
        response = requests.get(api_url, headers=headers)
        
        if response.status_code == 401:
            return None, None, "Error 401: No autorizado. Por favor, verifica tus credenciales."
        if response.status_code != 200:
            return None, None, f"Error al obtener información del modelo: Código de estado {response.status_code}"

        data = response.json()
        model_name = data.get('name', 'desconocido').replace(" ", "_")

        image_urls = []
        if 'modelVersions' in data and data['modelVersions']:
            model_version = data['modelVersions'][0]
            image_urls = [
                'https://civitai.com' + img['url'] if not img['url'].startswith('http') else img['url']
                for img in model_version.get('images', [])
            ]
            activation_text = ", ".join(model_version.get('trainedWords', []))
            file_size_kb = model_version.get('files', [{}])[0].get('sizeKb', "Unknown")
            description = self.html2text.handle(data.get("description", ""))

            model_info = {
                "name": model_name,
                "description": description,
                "tags": ", ".join(data.get("tags", [])),
                "creator": data.get("creator", {}).get("username", "Unknown"),
                "download count": data.get("stats", {}).get("downloadCount", "Unknown"),
                "favorite count": data.get("stats", {}).get("favoriteCount", "Unknown"),
                "comment count": data.get("stats", {}).get("commentCount", "Unknown"),
                "rating count": data.get("stats", {}).get("ratingCount", "Unknown"),
                "average rating": data.get("stats", {}).get("rating", "Unknown"),
                "activation text": activation_text,
                "file size (KB)": file_size_kb,
                "last updated": model_version.get('createdAt', "Unknown")
            }
            return model_info, image_urls, None
        else:
            return None, None, "No se encontraron versiones del modelo."

    def sanitize_filename(self, filename):
        return re.sub(r'[<>:"/\\|?*]', '_', filename)

    def download_file(self, url, dest_path, headers=None):
        with requests.get(url, headers=headers, stream=True) as r:
            r.raise_for_status()
            total_size = int(r.headers.get('content-length', 0))
            with open(dest_path, 'wb') as f:
                progress_bar = tqdm.tqdm(total=total_size, unit='B', unit_scale=True, unit_divisor=1024)
                for chunk in r.iter_content(chunk_size=8192):
                    if self.stop_download:  # Verificar si se debe detener la descarga
                        progress_bar.close()
                        return "Descarga detenida por el usuario."
                    f.write(chunk)
                    progress_bar.update(len(chunk))
                progress_bar.close()

    def download_model(self, url, download_folder, api_token, image_url=None):
        self.stop_download = False  # Reiniciar el estado de la variable de control
        model_info, image_urls, error_message = self.get_model_info(url, api_token)
        if error_message:
            return error_message

        download_path = os.path.join(directories[download_folder])
        if not os.path.exists(download_path):
            os.makedirs(download_path)

        model_id = self.extract_model_id(url)
        api_url = f'https://civitai.com/api/v1/models/{model_id}'
        headers = {'Authorization': f'Bearer {api_token}'} if api_token else {}
        response = requests.get(api_url, headers=headers)
        data = response.json()
        model_version = data['modelVersions'][0]
        file_info = model_version['files'][0]
        download_url = file_info['downloadUrl']
        download_url = 'https://civitai.com' + download_url if not download_url.startswith('http') else download_url

        try:
            model_file_name = f"{self.sanitize_filename(model_info['name'])}.safetensors"
            model_file_path = os.path.join(download_path, model_file_name)
            file_size = int(file_info.get('sizeKb', 0)) * 1024
            download_status = self.download_file(download_url, model_file_path, headers)
            
            if download_status == "Descarga detenida por el usuario.":
                return download_status

            if image_url:
                image_ext = os.path.splitext(image_url)[1]
                if image_ext.lower() not in ['.jpg', '.jpeg', '.png']:
                    image_ext = '.jpg'
                image_file_name = f"{self.sanitize_filename(model_info['name'])}{image_ext}"
                image_file_path = os.path.join(download_path, image_file_name)
                self.download_file(image_url, image_file_path, headers)
                image_status = f"Imagen descargada en: {image_file_path}"
            else:
                image_status = "No se seleccionó ninguna imagen de previsualización."

            json_file_name = f"{self.sanitize_filename(model_info['name'])}.json"
            json_file_path = os.path.join(download_path, json_file_name)
            with open(json_file_path, 'w') as json_file:
                json.dump(model_info, json_file, indent=4)

            return f"Descarga completada:\n- Modelo: {model_file_path}\n- {image_status}\n- JSON: {json_file_path}"

        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                return "Error 401: No autorizado. No se pudo descargar el modelo."
            else:
                return f"Error al descargar el modelo: {str(e)}"

    def stop_download_process(self):
        self.stop_download = True  # Cambiar la variable de control para detener la descarga

# Crear instancia de CivitaiDownloader
downloader = CivitaiDownloader()

def preview_model(url, api_token):
    model_info, image_urls, error_message = downloader.get_model_info(url, api_token)
    if error_message:
        return error_message, None

    # Generamos HTML estilizado para mostrar la información del modelo
    model_info_html = f"""
    <div style='font-family: Arial, sans-serif; padding: 15px; border-radius: 8px; background-color: #333; color: #f5f5f5; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); margin: auto; height: 300px; overflow-y: auto;'>
        <h2 style='color: #fff; border-bottom: 1px solid #444; padding-bottom: 10px; margin-bottom: 20px;'>{model_info['name']}</h2>
        <div style='margin-bottom: 15px;'>
            <p style='margin: 5px 0;'><strong style='color: #ddd;'>Description:</strong> {model_info['description']}</p>
            <p style='margin: 5px 0;'><strong style='color: #ddd;'>Tags:</strong> {model_info['tags']}</p>
            <p style='margin: 5px 0;'><strong style='color: #ddd;'>Creator:</strong> {model_info['creator']}</p>
            <p style='margin: 5px 0;'><strong style='color: #ddd;'>Download Count:</strong> {model_info['download count']}</p>
            <p style='margin: 5px 0;'><strong style='color: #ddd;'>Favorite Count:</strong> {model_info['favorite count']}</p>
            <p style='margin: 5px 0;'><strong style='color: #ddd;'>Comment Count:</strong> {model_info['comment count']}</p>
            <p style='margin: 5px 0;'><strong style='color: #ddd;'>Rating Count:</strong> {model_info['rating count']}</p>
            <p style='margin: 5px 0;'><strong style='color: #ddd;'>Average Rating:</strong> {model_info['average rating']}</p>
            <p style='margin: 5px 0;'><strong style='color: #ddd;'>Trigger Words:</strong> {model_info['activation text']}</p>
            <p style='margin: 5px 0;'><strong style='color: #ddd;'>Last Updated:</strong> {model_info['last updated']}</p>
        </div>
    </div>
    """

    return model_info_html, image_urls

def start_download(url, folder_key, api_token, image_url):
    return downloader.download_model(url, folder_key, api_token, image_url)

def stop_download():
    downloader.stop_download_process()
    return "Descarga detenida."

def image_selected(sel: gr.SelectData):
    # Asegúrate de que 'value' en sel contiene 'image' y 'url'
    if 'image' in sel.value and 'url' in sel.value['image']:
        return sel.value['image']['url']
    else:
        return "No se encontró la URL de la imagen."


with gr.Blocks() as interface:
    with gr.Tab("Downloader"):
        with gr.Row():
            with gr.Column(scale=2):
                with gr.Group():                   
                    link_input = gr.Textbox(label="Civitai URL", placeholder="Paste the Civitai link here")
                    api_token_input = gr.Textbox(label="API Token", placeholder="Enter your API token", value=api, visible=False)
                    folder_dropdown = gr.Dropdown(label="Download Folder", choices=list(directories.keys()), value="Flux Loras")
                    image_url_input = gr.Textbox(label="Image URL (optional)", placeholder="Enter the URL of the image or drag the image from the gallery", visible=True)
                    model_info_output = gr.HTML()
                    output_text = gr.Textbox(label="Download Status", lines=4)
        
                with gr.Column():
                    with gr.Row():
                        download_button = gr.Button("Download", variant="primary")
                        stop_button = gr.Button("Stop Download", variant="stop")       

            with gr.Column():
                with gr.Group():
                    model_image_gallery = gr.Gallery(label="Preview Images", columns=4, height=580)
         
                link_input.change(fn=preview_model, inputs=[link_input, api_token_input], outputs=[model_info_output, model_image_gallery])
                download_button.click(fn=start_download, inputs=[link_input, folder_dropdown, api_token_input, image_url_input], outputs=[output_text])
                stop_button.click(fn=stop_download, outputs=[output_text])  # Acción del botón de detener descarga
                model_image_gallery.select(image_selected, outputs=[image_url_input])

    with gr.TabItem("Update"):
        with gr.Row():
            with gr.Column(scale=1):
                model_folder_input = gr.Textbox(label="Carpeta de Modelos", placeholder="Introduce la ruta de la carpeta")
                check_button = gr.Button("Verificar Actualizaciones")
            
            with gr.Column(scale=4):
                output = gr.Dataframe(headers=["Modelo", "ID en Civitai", "Nombre del Modelo", "Versión del Modelo", "Error"], type="pandas", height=600)
        
        check_button.click(fn=check_for_updates, inputs=model_folder_input, outputs=output)

Script.interface = interface
script = Script()
script_callbacks.on_ui_tabs(script.on_ui_tabs)

Puedes mejorar este codigo?