#!/usr/bin/env python3
"""
test_catalyst_workflow.py — Catalyst Poster Generation via ComfyUI + Qwen-Image-2512
=====================================================================================

This script tests the Catalyst real-estate poster generation workflow by:
  1. Uploading a Ground Truth (architectural/floorplan) image to ComfyUI
  2. Uploading a Style Reference image (from Google/Pinterest) to ComfyUI
  3. Parsing raw UI input to extract marketing copy and aesthetic keywords
  4. Expanding the parsed input into a full Qwen-Image-2512-tuned prompt
  5. Dynamically injecting filenames and prompts into the workflow JSON
  6. Queuing the workflow on the ComfyUI server and polling for completion
  7. Downloading the final poster to a local output directory

Environment:
  - ComfyUI backend running Qwen-Image-2512 on AWS EC2 (4x NVIDIA L4, 96GB VRAM)
  - Model location: /home/ubuntu/models/Qwen-Image-2512 (diffusers sharded format)
  - ComfyUI location: /home/ubuntu/velocity/
  - Internal ComfyUI port: 8118 | External gateway port: 8288

Usage:
  python test_catalyst_workflow.py
"""

import os
import json
import re
import time
import base64
from pathlib import Path
from typing import Tuple, Optional

import requests
from PIL import Image

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# CONFIGURATION — Update COMFYUI_SERVER_URL with your AWS instance IP
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

COMFYUI_SERVER_URL: str = "http://<AWS_INSTANCE_IP>:8188"
"""
ComfyUI server URL. Replace <AWS_INSTANCE_IP> with the actual IP address.
  - For direct access (SSH tunnel): http://127.0.0.1:8118
  - For external access (if port 8188 is open): http://54.91.19.60:8188
  - Via Dream Weaver gateway (does NOT apply here): http://54.91.19.60:8288
Note: The internal ComfyUI port on the AWS instance is 8118. If SSH-tunnelling,
map local port 8188 to remote port 8118.
"""

INPUT_DIR: str = r"F:\Workin In Progress\DESINEURON\GITLAB\Project_Velocity\comfy_engine\test_inputs\Sagnik Test Sample New"
"""Base directory containing Ground Truth and Style Reference test images."""

OUTPUT_DIR: str = r"F:\Workin In Progress\DESINEURON\GITLAB\Project_Velocity\comfy_engine\test_outputs\Sagnik Test Sample New"
"""Directory to save generated poster outputs."""

WORKFLOW_JSON_PATH: str = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "..", "workflows", "catalyst_poster_qwen.json"
)
"""Path to the catalyst_poster_qwen.json workflow file (relative to this script)."""

# Node IDs in the workflow JSON (must match catalyst_poster_qwen.json)
NODE_ID_GROUND_TRUTH: str = "1"   # LoadImage node for Ground Truth
NODE_ID_STYLE_REF: str = "2"     # LoadImage node for Style Reference
NODE_ID_POSITIVE_PROMPT: str = "9"  # CLIPTextEncode node for positive prompt
NODE_ID_NEGATIVE_PROMPT: str = "10" # CLIPTextEncode node for negative prompt


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# FUNCTION 1: Prompt Parsing
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def process_prompt(raw_ui_input: str) -> Tuple[str, str]:
    """Parse raw UI input into aesthetic keywords and marketing copy.

    The raw input must contain marketing copy enclosed in double quotes.
    Everything outside the quotes is treated as aesthetic/style keywords.

    Args:
        raw_ui_input: Raw string from the UI, e.g.:
            'modern luxury warm lighting "Your Dream Home Awaits"'

    Returns:
        A tuple of (aesthetic_keywords, marketing_copy).

    Raises:
        ValueError: If no text enclosed in double quotes is found.

    Example:
        >>> process_prompt('art deco gold "Live in Elegance"')
        ('art deco gold', 'Live in Elegance')
    """
    match = re.search(r'"([^"]+)"', raw_ui_input)

    if not match:
        raise ValueError(
            "Marketing copy must be enclosed in double quotes. "
            "Example: 'modern luxury \"Your Dream Home Awaits\"'"
        )

    marketing_copy: str = match.group(1).strip()

    # Extract everything outside the quotes as aesthetic keywords
    aesthetic_keywords: str = raw_ui_input[:match.start()] + raw_ui_input[match.end():]
    aesthetic_keywords = re.sub(r'\s+', ' ', aesthetic_keywords).strip()

    return aesthetic_keywords, marketing_copy


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# FUNCTION 2: Prompt Expansion
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def expand_prompt(aesthetic_keywords: str, marketing_copy: str) -> str:
    """Expand parsed inputs into a full Qwen-Image-2512-optimized prompt.

    Constructs a semantically rich prompt formatted for Qwen-Image-2512's
    typography rendering capabilities. The prompt explicitly instructs the
    model to render text within the generated image.

    Args:
        aesthetic_keywords: Style descriptors (e.g., 'modern luxury warm lighting').
        marketing_copy: Exact text to render in the poster (e.g., 'Your Dream Home Awaits').

    Returns:
        A complete prompt string ready for CLIPTextEncode.

    Example:
        >>> expand_prompt('modern luxury', 'Live in Style')
        'A highly realistic, cinematic realestate marketing poster...'
    """
    return (
        f"A highly realistic, cinematic realestate marketing poster. "
        f"Interior style: {aesthetic_keywords}. "
        f"The image must prominently feature the exact text "
        f"'{marketing_copy}' written in elegant, modern, highly legible "
        f"typography. Professional lighting, 8k resolution, photorealistic "
        f"quality, detailed textures."
    )


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# FUNCTION 3: Image Upload
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def upload_image(image_path: str) -> str:
    """Upload an image to the ComfyUI server for use in workflows.

    Opens the image, converts to RGB if necessary, saves as a
    temporary PNG, and uploads via the /upload/image endpoint.

    Args:
        image_path: Absolute path to the image file on disk.

    Returns:
        The server-side filename assigned by ComfyUI (used in workflow JSON).

    Raises:
        FileNotFoundError: If the image file does not exist.
        requests.exceptions.ConnectionError: If the ComfyUI server is unreachable.
        requests.exceptions.Timeout: If the upload times out.
        RuntimeError: If the server returns an unexpected response.
    """
    path = Path(image_path)
    if not path.exists():
        raise FileNotFoundError(f"Image not found: {image_path}")

    # Open and ensure RGB mode
    img = Image.open(path)
    if img.mode != "RGB":
        img = img.convert("RGB")

    # Save to a temporary PNG buffer for upload
    import tempfile
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        tmp_path = tmp.name
        img.save(tmp_path, format="PNG")

    try:
        with open(tmp_path, "rb") as f:
            files = {
                "image": (path.name, f, "image/png")
            }
            data = {
                "overwrite": "true"
            }
            response = requests.post(
                f"{COMFYUI_SERVER_URL}/upload/image",
                files=files,
                data=data,
                timeout=60
            )
            response.raise_for_status()

        result = response.json()
        server_filename: str = result.get("name", "")

        if not server_filename:
            raise RuntimeError(
                f"ComfyUI upload returned unexpected response: {result}"
            )

        print(f"  ✓ Uploaded '{path.name}' → server filename: '{server_filename}'")
        return server_filename

    finally:
        # Clean up temp file
        try:
            os.unlink(tmp_path)
        except OSError:
            pass


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# FUNCTION 4: Execute Workflow
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def execute_workflow(
    workflow_json: dict,
    prompt_text: str,
    ground_truth_filename: str,
    style_ref_filename: str
) -> str:
    """Inject dynamic values into the workflow JSON and queue it on ComfyUI.

    Updates the following nodes in the workflow:
      - Node 1 (LoadImage): Sets Ground Truth filename
      - Node 2 (LoadImage): Sets Style Reference filename
      - Node 9 (CLIPTextEncode): Sets the expanded positive prompt

    Args:
        workflow_json: The loaded workflow JSON dict (API format).
        prompt_text: The expanded prompt string from expand_prompt().
        ground_truth_filename: Server-side filename of the ground truth image.
        style_ref_filename: Server-side filename of the style reference image.

    Returns:
        The prompt_id string from ComfyUI's queue response.

    Raises:
        requests.exceptions.ConnectionError: If the server is unreachable.
        requests.exceptions.Timeout: If the request times out.
        KeyError: If expected node IDs are missing from the workflow JSON.
    """
    # Deep copy to avoid mutating the original
    import copy
    wf = copy.deepcopy(workflow_json)

    # Inject Ground Truth image filename
    wf[NODE_ID_GROUND_TRUTH]["inputs"]["image"] = ground_truth_filename

    # Inject Style Reference image filename
    wf[NODE_ID_STYLE_REF]["inputs"]["image"] = style_ref_filename

    # Inject expanded positive prompt
    wf[NODE_ID_POSITIVE_PROMPT]["inputs"]["text"] = prompt_text

    # Build the API payload
    payload = {
        "prompt": wf,
        "client_id": f"catalyst_{int(time.time())}"
    }

    print(f"  → Queuing workflow on {COMFYUI_SERVER_URL}/prompt ...")

    response = requests.post(
        f"{COMFYUI_SERVER_URL}/prompt",
        json=payload,
        timeout=30
    )
    response.raise_for_status()

    result = response.json()
    prompt_id: str = result.get("prompt_id", "")

    if not prompt_id:
        raise RuntimeError(
            f"ComfyUI /prompt returned unexpected response: {result}"
        )

    print(f"  ✓ Queued successfully. prompt_id: {prompt_id}")
    return prompt_id


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# FUNCTION 5: Poll for Completion
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def wait_for_completion(
    prompt_id: str,
    timeout: int = 300,
    poll_interval: int = 2
) -> dict:
    """Poll the ComfyUI history endpoint until the workflow completes.

    Repeatedly checks /history/{prompt_id} for output images. The Qwen-Image-2512
    model with 50 inference steps on 4x L4 GPUs typically takes 30-90 seconds.

    Args:
        prompt_id: The prompt ID returned by execute_workflow().
        timeout: Maximum seconds to wait before raising TimeoutError.
        poll_interval: Seconds between poll requests.

    Returns:
        The history dict for this prompt_id (contains output image metadata).

    Raises:
        TimeoutError: If the workflow doesn't complete within timeout seconds.
        requests.exceptions.ConnectionError: If the server is unreachable.
        RuntimeError: If the workflow reports an error status.
    """
    history_url = f"{COMFYUI_SERVER_URL}/history/{prompt_id}"
    start_time = time.time()
    poll_count = 0

    print(f"  ⏳ Polling for completion (timeout: {timeout}s) ...")

    while time.time() - start_time < timeout:
        time.sleep(poll_interval)
        poll_count += 1

        try:
            response = requests.get(history_url, timeout=10)

            if response.status_code == 200:
                history = response.json()
                prompt_history = history.get(prompt_id, {})

                # Check for error status
                status_info = prompt_history.get("status", {})
                if status_info.get("status_str") == "error":
                    error_msgs = status_info.get("messages", ["Unknown error"])
                    raise RuntimeError(
                        f"Workflow execution failed: {error_msgs}"
                    )

                # Check for output images
                outputs = prompt_history.get("outputs", {})
                for node_id, node_output in outputs.items():
                    if "images" in node_output and node_output["images"]:
                        elapsed = time.time() - start_time
                        print(
                            f"  ✓ Completed in {elapsed:.1f}s "
                            f"({poll_count} polls)"
                        )
                        return prompt_history

        except requests.exceptions.ConnectionError:
            # Server might be busy with GPU inference, retry
            print(f"    Poll #{poll_count}: Connection interrupted, retrying...")
        except requests.exceptions.Timeout:
            print(f"    Poll #{poll_count}: Timeout, retrying...")

    raise TimeoutError(
        f"Workflow did not complete within {timeout} seconds "
        f"(prompt_id: {prompt_id})"
    )


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# FUNCTION 6: Download Output
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def download_output(history: dict, output_dir: str) -> str:
    """Extract and download the generated poster from ComfyUI history.

    Reads the output image metadata from the history response, constructs
    the /view URL, downloads the image, and saves it with a timestamped
    filename.

    Args:
        history: The prompt history dict returned by wait_for_completion().
        output_dir: Local directory to save the output image.

    Returns:
        The absolute path to the saved output image.

    Raises:
        RuntimeError: If no output images are found in the history.
        requests.exceptions.ConnectionError: If the server is unreachable.
    """
    # Find the output image in the history
    output_image: Optional[dict] = None

    for node_id, node_output in history.get("outputs", {}).items():
        images = node_output.get("images", [])
        if images:
            output_image = images[0]
            break

    if not output_image:
        raise RuntimeError("No output images found in workflow history")

    # Construct the ComfyUI /view URL
    filename = output_image["filename"]
    subfolder = output_image.get("subfolder", "")
    img_type = output_image.get("type", "output")

    view_url = (
        f"{COMFYUI_SERVER_URL}/view"
        f"?filename={filename}"
        f"&subfolder={subfolder}"
        f"&type={img_type}"
    )

    print(f"  ⬇ Downloading: {filename} ...")

    response = requests.get(view_url, stream=True, timeout=60)
    response.raise_for_status()

    # Save with timestamp
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    output_filename = f"catalyst_poster_{timestamp}.png"
    output_path = os.path.join(output_dir, output_filename)

    with open(output_path, "wb") as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)

    file_size_mb = os.path.getsize(output_path) / (1024 * 1024)
    print(f"  ✓ Saved: {output_path} ({file_size_mb:.1f} MB)")

    return output_path


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# MAIN EXECUTION
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

if __name__ == "__main__":
    print("=" * 72)
    print("  CATALYST POSTER GENERATION — Qwen-Image-2512 Workflow Test")
    print("=" * 72)

    # ── Ensure output directory exists ──
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    print(f"\n📂 Output dir: {OUTPUT_DIR}")

    # ── Load workflow JSON ──
    workflow_path = os.path.normpath(WORKFLOW_JSON_PATH)
    print(f"📄 Workflow:   {workflow_path}")

    try:
        with open(workflow_path, "r", encoding="utf-8") as f:
            workflow = json.load(f)
        print(f"  ✓ Loaded workflow ({len(workflow)} nodes)")
    except FileNotFoundError:
        print(f"  ✗ ERROR: Workflow file not found: {workflow_path}")
        print("    Ensure catalyst_poster_qwen.json is in ../workflows/")
        exit(1)
    except json.JSONDecodeError as e:
        print(f"  ✗ ERROR: Invalid JSON in workflow file: {e}")
        exit(1)

    # ── Define test inputs ──
    # Update these paths to your actual test images
    ground_truth_image = os.path.join(INPUT_DIR, "ground_truth.jpg")
    style_reference_image = os.path.join(INPUT_DIR, "style_reference.jpg")

    raw_prompt = (
        'modern luxury warm ambient lighting premium materials '
        'golden hour cinematic architectural photography '
        '"Your Dream Home Awaits"'
    )

    print(f"\n🖼️  Ground Truth:    {ground_truth_image}")
    print(f"🎨 Style Reference: {style_reference_image}")
    print(f"📝 Raw Prompt:      {raw_prompt}")

    # ── Step 1: Parse the prompt ──
    print("\n── Step 1: Parsing prompt ──")
    try:
        aesthetic_keywords, marketing_copy = process_prompt(raw_prompt)
        print(f"  Keywords: {aesthetic_keywords}")
        print(f"  Copy:     \"{marketing_copy}\"")
    except ValueError as e:
        print(f"  ✗ ERROR: {e}")
        exit(1)

    # ── Step 2: Expand the prompt ──
    print("\n── Step 2: Expanding prompt ──")
    expanded = expand_prompt(aesthetic_keywords, marketing_copy)
    print(f"  Expanded ({len(expanded)} chars):")
    print(f"  {expanded[:120]}...")

    # ── Step 3: Upload images ──
    print("\n── Step 3: Uploading images ──")
    try:
        gt_filename = upload_image(ground_truth_image)
        sr_filename = upload_image(style_reference_image)
    except FileNotFoundError as e:
        print(f"  ✗ ERROR: {e}")
        print("    Place test images in the INPUT_DIR directory.")
        exit(1)
    except requests.exceptions.ConnectionError as e:
        print(f"  ✗ CONNECTION ERROR: Cannot reach {COMFYUI_SERVER_URL}")
        print(f"    Details: {e}")
        print("    Ensure ComfyUI is running and the URL is correct.")
        print("    If using SSH tunnel: ssh -L 8188:127.0.0.1:8118 ...")
        exit(1)
    except requests.exceptions.Timeout:
        print(f"  ✗ TIMEOUT: Upload timed out to {COMFYUI_SERVER_URL}")
        exit(1)
    except Exception as e:
        print(f"  ✗ UNEXPECTED ERROR during upload: {e}")
        exit(1)

    # ── Step 4: Execute workflow ──
    print("\n── Step 4: Executing workflow ──")
    try:
        prompt_id = execute_workflow(
            workflow_json=workflow,
            prompt_text=expanded,
            ground_truth_filename=gt_filename,
            style_ref_filename=sr_filename
        )
    except requests.exceptions.ConnectionError as e:
        print(f"  ✗ CONNECTION ERROR: {e}")
        exit(1)
    except requests.exceptions.Timeout:
        print(f"  ✗ TIMEOUT: Could not queue workflow")
        exit(1)
    except KeyError as e:
        print(f"  ✗ WORKFLOW ERROR: Missing node ID {e} in workflow JSON")
        exit(1)
    except Exception as e:
        print(f"  ✗ UNEXPECTED ERROR: {e}")
        exit(1)

    # ── Step 5: Poll for completion ──
    print("\n── Step 5: Waiting for completion ──")
    try:
        history = wait_for_completion(
            prompt_id=prompt_id,
            timeout=300,
            poll_interval=2
        )
    except TimeoutError as e:
        print(f"  ✗ TIMEOUT: {e}")
        exit(1)
    except RuntimeError as e:
        print(f"  ✗ EXECUTION ERROR: {e}")
        exit(1)
    except Exception as e:
        print(f"  ✗ UNEXPECTED ERROR: {e}")
        exit(1)

    # ── Step 6: Download output ──
    print("\n── Step 6: Downloading output ──")
    try:
        output_path = download_output(
            history=history,
            output_dir=OUTPUT_DIR
        )
    except RuntimeError as e:
        print(f"  ✗ ERROR: {e}")
        exit(1)
    except requests.exceptions.ConnectionError as e:
        print(f"  ✗ DOWNLOAD ERROR: {e}")
        exit(1)
    except Exception as e:
        print(f"  ✗ UNEXPECTED ERROR: {e}")
        exit(1)

    # ── Success ──
    print("\n" + "=" * 72)
    print(f"  ✅ SUCCESS — Poster saved to:")
    print(f"     {output_path}")
    print("=" * 72)