import os
import json
import requests
import random
import uuid
import websocket as ws_client
from dotenv import load_dotenv
from urllib.request import urlopen, Request
from urllib.parse import urlencode

# Load environment variables
load_dotenv()

server_address = os.getenv('COMFYUI_SERVER_ADDRESS', 'localhost:8188')
client_id = str(uuid.uuid4())
ollama_server_address = os.getenv('OLLAMA_SERVER_ADDRESS', 'localhost:11434')
ollama_server_url = f"http://{ollama_server_address}/v1/chat/completions"
ollama_model = os.getenv('OLLAMA_MODEL', 'llama3.1:latest')  # Load model from .env
system_prompt = """
You are an AI model that transforms user input into concise, creative, and visually descriptive prompts for AI image generation. If the user input is empty, generate a creative, detailed description automatically. Always respond with a single, visually descriptive line, incorporating elements such as hyperrealistic, 4K, and detailed imagery when applicable.

Guidelines:
1. **Be Specific and Clear**: Ensure the response describes the subject, setting, and style in a visually engaging manner.
2. **Keep It Brief**: Limit responses to one or two sentences that provide a vivid image while remaining concise.
3. **Encourage Creativity**: Incorporate artistic styles like hyperrealism, 4K resolution, cinematic lighting, or surrealism where appropriate.
4. **Handle Empty Input**: If the input is empty, create an imaginative and detailed prompt using the aforementioned elements.
5. **No Extra Communication**: Respond only with the transformed prompt without additional commentary.

### Example Inputs and Outputs:

**Input**: "A futuristic city at night with neon lights and flying cars."
**Output**: "A neon-lit futuristic cityscape at night, with flying cars darting between towering skyscrapers, glowing in vibrant 4K detail."

**Input**: "A serene beach with palm trees and a sunset."
**Output**: "A serene beach at sunset, palm trees swaying gently as the golden light reflects off the 4K ocean waves."

**Input**: "A magical forest with glowing trees and animals."
**Output**: "A hyperrealistic forest where glowing trees light the way for mystical creatures under a twilight sky, captured in stunning 4K."

**Input**: ""
**Output**: "An intricately detailed 4K portrait of a warrior in golden armor, standing proudly under the dramatic glow of cinematic lighting."

**Input**: "A watercolor painting of a rainy street."
**Output**: "A dreamy watercolor painting of a rainy street, soft reflections glistening in the puddles under warm streetlights."
"""
# Service to get quick prompts data
def get_quick_prompts_data():
    with open("quick_prompts.json", "r") as f:
        return json.load(f)

# Service to ask LLM for response
def ask_llm_service(positive_prompt: str):
    ollama_request = {
        "model": ollama_model,  # Use model from .env
        "messages": [
            {
                "role": "system",
                "content": system_prompt
            },
            {
                "role": "user",
                "content": positive_prompt
            }
        ]
    }

    response = requests.post(ollama_server_url, json=ollama_request, headers={"Content-Type": "application/json"})
    response_data = response.json()
    return response_data["choices"][0]["message"]["content"]

# Service to queue a prompt
def queue_prompt(prompt):
    p = {"prompt": prompt, "client_id": client_id}
    data = json.dumps(p).encode('utf-8')
    req = Request(f"http://{server_address}/prompt", data=data)
    return json.loads(urlopen(req).read())

# Service to get image
def get_image(filename, subfolder, folder_type):
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    url_values = urlencode(data)
    with urlopen(f"http://{server_address}/view?{url_values}") as response:
        return response.read()

# Service to get history
def get_history(prompt_id):
    with urlopen(f"http://{server_address}/history/{prompt_id}") as response:
        return json.loads(response.read())

# WebSocket image generation service
async def get_images(ws, prompt):
    prompt_id = queue_prompt(prompt)['prompt_id']
    output_images = {}

    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message['type'] == 'executing':
                data = message['data']
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    break
        else:
            continue

    history = get_history(prompt_id)[prompt_id]
    for o in history['outputs']:
        for node_id in history['outputs']:
            node_output = history['outputs'][node_id]
            if 'images' in node_output:
                images_output = []
                for image in node_output['images']:
                    image_data = get_image(image['filename'], image['subfolder'], image['type'])
                    images_output.append(image_data)
                output_images[node_id] = images_output

    return output_images

# Main image generation function
async def generate_images(positive_prompt, negative_prompt, steps=25, resolution=(512, 512)):
    ws = ws_client.WebSocket()
    ws.connect(f"ws://{server_address}/ws?clientId={client_id}")

    with open("workflow.json", "r", encoding="utf-8") as f:
        workflow_data = f.read()

    workflow = json.loads(workflow_data)

    workflow["6"]["inputs"]["text"] = positive_prompt
    workflow["7"]["inputs"]["text"] = negative_prompt
    workflow["3"]["inputs"]["steps"] = steps
    workflow["5"]["inputs"]["width"] = resolution[0]
    workflow["5"]["inputs"]["height"] = resolution[1]
    seed = random.randint(1, 1000000000)
    workflow["3"]["inputs"]["seed"] = seed

    images = await get_images(ws, workflow)

    ws.close()

    return images, seed