import os
import sys
import zipfile
import gradio as gr
import numpy as np
from PIL import Image, PngImagePlugin
import asyncio
import io
import json
import re
import random
import base64

# Unzip the novelai_api package if it's not already extracted
if not os.path.exists("novelai_api"):
    with zipfile.ZipFile("novelai_api.zip", "r") as zip_ref:
        zip_ref.extractall(".")

# Add the extracted directory to Python path
sys.path.append(os.path.abspath("novelai_api"))

# Import NovelAIAPI
from novelai_api import NovelAIAPI
from novelai_api.ImagePreset import ImageModel, ImageGenerationType

# Retrieve the NovelAI API token from environment variables
ACCESS_TOKEN = os.getenv("NOVELAI_ACCESS_TOKEN")
if ACCESS_TOKEN is None:
    raise ValueError("API token not found. Please set it as a secret in Hugging Face settings.")

# Paths to JSON data files
ARTISTS_PATH = 'artists.json'
CHARACTERS_PATH = 'characters.json'

# Load data from a JSON file
def load_json_file(file_path):
    try:
        with open(file_path, 'r') as file:
            return json.load(file)
    except Exception as e:
        print(f"Error reading {file_path}: {e}")
        return []

# Parse metadata from an image file
def parse_image_metadata(image_path):
    params = {}
    try:
        with Image.open(image_path) as img:
            if isinstance(img.info, PngImagePlugin.PngInfo):
                metadata = img.info.get("parameters")
                if metadata:
                    metadata_json = json.loads(metadata)
                    params['prompt'] = metadata_json.get("prompt", "")
                    params['steps'] = metadata_json.get("steps", 28)
                    params['height'] = metadata_json.get("height", 832)
                    params['width'] = metadata_json.get("width", 1216)
                    params['scale'] = metadata_json.get("scale", 6.5)
                    params['uc'] = metadata_json.get("uc", "")
                    params['sampler'] = metadata_json.get("sampler", "k_euler_ancestral")
                    params['noise_schedule'] = metadata_json.get("noise_schedule", "karras")
                    params['controlnet_strength'] = metadata_json.get("controlnet_strength", 1) == 1
                    params['prefer_brownian'] = metadata_json.get("prefer_brownian", True)
                    params['cfg_rescale'] = metadata_json.get("cfg_rescale", 0.2)
                    params['seed'] = metadata_json.get("seed", random.randint(0, 2**32 - 1))  # Default to a random seed if not found
                    params['character_name'] = extract_character_from_prompt(params['prompt'])
    except (json.JSONDecodeError, KeyError):
        print("Error parsing metadata JSON.")
    except Exception as e:
        print(f"Error reading image metadata: {e}")
    return params

# Extract character from the beginning of the positive prompt
def extract_character_from_prompt(prompt):
    return prompt.split(",")[0].strip() if prompt else ""

# Function to get random characters
def get_random_characters(count):
    characters = load_json_file(CHARACTERS_PATH)
    return ', '.join(random.sample([char['character_name'] for char in characters], count))

# Function to get random artists
def get_random_artists(count):
    artists = load_json_file(ARTISTS_PATH)
    return ', '.join(random.sample([artist['artist_name'] for artist in artists], count))

# Function to parse and replace placeholders in the prompt
def parse_prompt(prompt):
    # Replace randChar(n) with n random characters
    prompt = re.sub(r'randChar\((\d+)\)', lambda m: get_random_characters(int(m.group(1))), prompt)
    # Replace randArtist(n) with n random artists
    prompt = re.sub(r'randArtist\((\d+)\)', lambda m: get_random_artists(int(m.group(1))), prompt)
    return prompt

# Generate image based on prompt and metadata parameters
async def generate_image(api, metadata_params):
    parameters = {
        "negative_prompt": metadata_params.get("uc", ""),
        "height": metadata_params.get("height", 832),
        "width": metadata_params.get("width", 1216),
        "scale": metadata_params.get("scale", 6.5),
        "steps": metadata_params.get("steps", 28),
        "sampler": metadata_params.get("sampler", "k_euler_ancestral"),
        "noise_schedule": metadata_params.get("noise_schedule", "karras"),
        "controlnet_strength": 1.0 if metadata_params.get("controlnet_strength", True) else 0.0,
        "prefer_brownian": metadata_params.get("prefer_brownian", True),
        "cfg_rescale": metadata_params.get("cfg_rescale", 0.2),
        "seed": metadata_params.get("seed", random.randint(0, 2**32 - 1)),
    }

    prompt = metadata_params.get("prompt") or "default prompt"  # Set a default prompt if not in metadata
    prompt = parse_prompt(prompt)  # Parse the prompt to replace placeholders

    try:
        async for _, img in api.low_level.generate_image(prompt, ImageModel.Anime_v3, ImageGenerationType.NORMAL, parameters):
            return img  # Returning raw bytes
    except Exception as e:
        print(f"An error occurred during image generation: {e}")
        return None

# Gradio Interface with Blocks
with gr.Blocks(css="#col-container { max-width: 800px; }") as demo:
    characters = load_json_file(CHARACTERS_PATH)
    character_options = ['random'] + [char['character_name'] for char in characters]
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown("<div id='title'>NovelAI Image Generator</div>")

        character_selection = gr.Dropdown(character_options, label="Enter or Select Character", value="random")
        additional_tags = gr.Textbox(label="Additional Tags (comma-separated)", placeholder="e.g., nsfw, vibrant colors")
        artists = gr.Textbox(label="Artists (comma-separated) or leave empty for random", placeholder="e.g., artist1, artist2")
        quality_tags = gr.Checkbox(label="Include Quality Tags", value=True)
        negative_prompt = gr.Textbox(label="Negative Prompt", value="{{worst quality, bad quality, censored}}, amputee, deformed")

        with gr.Accordion("Advanced Options", open=False):
            width = gr.Slider(512, 1216, step=64, label="Width", value=832)
            height = gr.Slider(512, 1216, step=64, label="Height", value=1216)
            scale = gr.Slider(1, 100, step=0.01, label="Scale", value=6.5)
            steps = gr.Slider(1, 50, step=1, label="Steps", value=28)
            sampler = gr.Dropdown(["k_euler", "k_euler_ancestral", "ddim", "k_dpm_2"], label="Sampler", value="k_euler_ancestral")
            noise_schedule = gr.Dropdown(["karras", "polyexponential"], label="Noise Schedule", value="karras")
            controlnet_strength = gr.Checkbox(label="ControlNet Strength", value=True)
            prefer_brownian = gr.Checkbox(label="Prefer Brownian Noise", value=True)
            cfg_rescale = gr.Slider(0, 1, step=0.01, label="CFG Rescale", value=0.2)

        metadata_upload = gr.Image(label="Upload Image with Metadata", type="filepath")

        # Function to handle image upload and extract metadata
        def update_parameters_with_metadata(image_path):
            parsed_params = parse_image_metadata(image_path)
            character_name = parsed_params.get('character_name', 'random')  # Extracted character
            return (
                character_name,
                parsed_params.get('uc', ''),
                parsed_params.get('width', 832),
                parsed_params.get('height', 1216),
                parsed_params.get('scale', 6.5),
                parsed_params.get('steps', 28),
                parsed_params.get('sampler', 'k_euler_ancestral'),
                parsed_params.get('noise_schedule', 'karras'),
                parsed_params.get('controlnet_strength', True),
                parsed_params.get('prefer_brownian', True),
                parsed_params.get('cfg_rescale', 0.2)
            )

        metadata_upload.change(
            fn=update_parameters_with_metadata,
            inputs=metadata_upload,
            outputs=[character_selection, negative_prompt, width, height, scale, steps, sampler, noise_schedule, controlnet_strength, prefer_brownian, cfg_rescale]
        )

        generate_button = gr.Button("Generate Image")
        result_image = gr.Image(type="numpy", label="Generated Image")
        result_message = gr.Textbox(label="Status")

    generate_button.click(
        fn=lambda *args: asyncio.run(generate_image(*args)),
        inputs=[
            character_selection, additional_tags, artists, quality_tags, negative_prompt,
            width, height, scale, steps, sampler, noise_schedule, controlnet_strength,
            prefer_brownian, cfg_rescale
        ],
        outputs=[result_image, result_message],
    )

if __name__ == "__main__":
    demo.launch()