import gradio as gr
import torch
import json
import time
import subprocess
import os
from transformers import AutoTokenizer, AutoModelForCausalLM

# Initialize DeepSeek
model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.to("cpu")

def generate_prompts_and_image(input_text):
    start_time = time.time()

    # Generate prompts with DeepSeek
    prompt = f"""
    Input: "{input_text}"
    Task: Generate concise 'Positive' and 'Negative' AI image prompts for Stable Diffusion based on the input above. Output the prompts directly, no extra text or examples.
    """
    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
    outputs = model.generate(**inputs, max_new_tokens=50, temperature=0.7, top_p=0.9, do_sample=True)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
    if response.startswith(prompt):
        response = response[len(prompt):].strip()

    # Split response into positive and negative (assuming two lines)
    lines = [line.strip() for line in response.split("\n") if line.strip()]
    positive = lines[0] if lines else "No positive prompt generated"
    negative = lines[1] if len(lines) > 1 else "No negative prompt generated"

    # Debug: Check if ComfyUI and model exist
    comfyui_dir = os.path.join(os.path.dirname(__file__), "ComfyUI")
    checkpoint_path = os.path.join(comfyui_dir, "models", "checkpoints", "v1-5-pruned-emaonly-fp16.safetensors")
    if not os.path.exists(comfyui_dir):
        return {"error": f"ComfyUI directory not found at {comfyui_dir}", "time_taken": f"{time.time() - start_time:.2f} seconds"}
    if not os.path.exists(checkpoint_path):
        return {"error": f"Checkpoint not found at {checkpoint_path}", "time_taken": f"{time.time() - start_time:.2f} seconds"}

    # Load and debug ComfyUI workflow
    workflow_path = "workflow.json"
    temp_workflow_path = "temp_workflow.json"
    if not os.path.exists(workflow_path):
        return {"error": "workflow.json not found", "time_taken": f"{time.time() - start_time:.2f} seconds"}

    with open(workflow_path, "r") as f:
        workflow = json.load(f)

    # Debug: Print workflow nodes to verify IDs
    print("Workflow nodes:", json.dumps(workflow["nodes"], indent=2))

    # Use hardcoded node IDs from workflow.json (6 for positive, 7 for negative)
    positive_id, negative_id = 6, 7  # Hardcoded based on your workflow.json
    try:
        if str(positive_id) not in workflow or str(negative_id) not in workflow:
            return {"error": f"Node IDs {positive_id} or {negative_id} not found in workflow.json", "time_taken": f"{time.time() - start_time:.2f} seconds"}
        
        # Inject prompts into nodes 6 (positive) and 7 (negative)
        workflow[str(positive_id)]["widgets_values"][0] = positive  # Positive prompt (CLIPTextEncode)
        workflow[str(negative_id)]["widgets_values"][0] = negative  # Negative prompt (CLIPTextEncode)
    except KeyError as e:
        return {"error": f"Error injecting prompts: {str(e)}", "time_taken": f"{time.time() - start_time:.2f} seconds"}

    # Save temporary workflow
    with open(temp_workflow_path, "w") as f:
        json.dump(workflow, f)

    # Run ComfyUI via subprocess
    comfyui_main = os.path.join(comfyui_dir, "main.py")
    output_dir = os.path.join(comfyui_dir, "output")  # ComfyUI default output dir
    os.makedirs(output_dir, exist_ok=True)

    try:
        # Note: ComfyUI's main.py doesn't use --input-directory or --output-directory directly; it reads from workflow
        result = subprocess.run(
            ["python", comfyui_main],
            cwd=comfyui_dir,
            capture_output=True,
            text=True,
            check=True
        )
        # ComfyUI saves images as ComfyUI_<timestamp>.png in output/
        image_files = [f for f in os.listdir(output_dir) if f.startswith("ComfyUI") and f.endswith(".png")]
        if not image_files:
            image = f"ComfyUI ran but no image found: {result.stdout}\n{result.stderr}"
        else:
            image = os.path.join(output_dir, image_files[-1])  # Use the latest image
    except subprocess.CalledProcessError as e:
        image = f"ComfyUI failed: {e.stdout}\n{e.stderr}"

    elapsed_time = time.time() - start_time
    return {
        "positive": positive,
        "negative": negative,
        "image": image,
        "time_taken": f"{elapsed_time:.2f} seconds"
    }

def gradio_interface(input_text):
    result = generate_prompts_and_image(input_text)
    return result["image"], json.dumps({
        "positive": result["positive"],
        "negative": result["negative"],
        "time_taken": result["time_taken"]
    }, indent=2)

demo = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(label="Input Text", lines=10, placeholder="Paste your input text here..."),
    outputs=[
        gr.Image(label="Generated Image"),
        gr.Textbox(label="Generated Prompts (JSON)")
    ],
    title="Prompt and Image Generator"
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)