import os
import requests
import time
import json
from pathlib import Path

# Config
GATEWAY_URL = "http://54.91.19.60:8082" # Active IP
INPUT_DIR = Path(r"f:\Workin In Progress\DESINEURON\GITLAB\Project_Velocity\comfy_engine\test_inputs\Abantika Test Sample")
OUTPUT_DIR = Path(r"f:\Workin In Progress\DESINEURON\GITLAB\Project_Velocity\comfy_engine\test_outputs\Abantika Test Samples - Test 2")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# 4 Styles requested by user
STYLES = [
    {
        "id": "gothic_industrial",
        "keywords": "Gothic, Indusrial Design, Black, Wood Textures and Accents, Bare Metal Edison Bulbs, Red Silk",
        "denoise": 0.75
    },
    {
        "id": "greek_minimal",
        "keywords": "Greek Aesthetic, Minimal, Grand, White, Deep Blue Silk, Emral Green Marbel",
        "denoise": 0.75
    },
    {
        "id": "turkish_vintage",
        "keywords": "Turkish Interioir, Mosaic Work, Vintage, Intricate Work, Royal",
        "denoise": 0.75
    },
    {
        "id": "bali_modern",
        "keywords": "Bali Aesthetic, Modern, Minimal, Stone, Live Indoor House Plants, Indonesian Suar wood (Samanea saman)",
        "denoise": 0.75
    }
]

def map_filename_to_room_type(filename: str) -> str:
    name = filename.lower()
    if "bed-room" in name or "bedroom" in name or "bed" in name:
        return "bedroom"
    elif "bath-room" in name or "bathroom" in name or "bath" in name:
        return "bathroom"
    elif "kitchen" in name:
        return "kitchen"
    elif "dining" in name:
        return "dining_room"
    elif "balcony" in name:
        return "balcony"
    return "living_room"  # default

def run_job(image_path: Path, style_cfg: dict):
    room_type = map_filename_to_room_type(image_path.name)
    style_id = style_cfg["id"]
    keywords = style_cfg["keywords"]
    
    print(f"\n--- Processing {image_path.name} with style {style_id.upper()} ({room_type}) ---")
    
    # Check if already processed
    img_out = OUTPUT_DIR / f"{image_path.stem}_{style_id}.png"
    if img_out.exists():
        print(f"[{style_id}] Already processed {img_out.name}, skipping.")
        return

    # 1. Start generation
    try:
        with open(image_path, "rb") as f:
            files = {"image": (image_path.name, f, "image/jpeg")}
            data = {
                "keywords": keywords,
                "room_type": room_type,
                "denoise": style_cfg["denoise"]
            }
            res = requests.post(f"{GATEWAY_URL}/dream-weaver", files=files, data=data, timeout=180)
            res.raise_for_status()
            
            job_data = res.json()
            job_id = job_data["job_id"]
            prompt_preview = job_data.get("prompt_preview", "")
            print(f"[{style_id}] Job submitted: {job_id}")
            print(f"[{style_id}] Prompt Preview: {prompt_preview[:150]}...")
            
    except requests.exceptions.RequestException as e:
        print(f"Failed to submit {image_path.name}: {e}")
        return

    # 2. Poll status
    status_url = f"{GATEWAY_URL}/dream-weaver/status/{job_id}"
    ready = False
    
    # Poll for up to 6 minutes
    s_data = {}
    for i in range(180):
        time.sleep(2)
        try:
            s_res = requests.get(status_url, timeout=10)
            if s_res.status_code == 200:
                s_data = s_res.json()
                if s_data.get("ready"):
                    ready = True
                    break
                elif s_data.get("status") == "error":
                    print(f"Job failed: {s_data.get('error')}")
                    return
        except requests.exceptions.RequestException as e:
            print(f"Polling error: {e}")
            
    if not ready:
        print(f"Timeout waiting for job {job_id}")
        return
        
    # 3. Download image and save prompt
    print(f"[{style_id}] Job ready, downloading...")
    
    # Save prompt data locally
    prompt_file = OUTPUT_DIR / f"{image_path.stem}_{style_id}_prompt.json"
    with open(prompt_file, "w") as pf:
        json.dump(s_data, pf, indent=2)
        
    result_url = f"{GATEWAY_URL}/dream-weaver/result/{job_id}"
    try:
        r_res = requests.get(result_url, stream=True, timeout=60)
        r_res.raise_for_status()
        
        with open(img_out, "wb") as f:
            for chunk in r_res.iter_content(chunk_size=8192):
                f.write(chunk)
                
        print(f"[{style_id}] Saved {img_out.name}")
    except requests.exceptions.RequestException as e:
        print(f"Failed to download result for {job_id}: {e}")

def main():
    # Wait for gateway and Ollama to be fully ready
    print(f"Checking Gateway at {GATEWAY_URL}/health...")
    for _ in range(3):
        try:
            h = requests.get(f"{GATEWAY_URL}/health", timeout=5)
            if h.status_code == 200:
                print("Gateway is UP.")
                break
        except requests.exceptions.RequestException:
            print("Waiting for gateway...")
            time.sleep(5)
    else:
        print("Gateway failed to answer health check. Exiting.")
        return

    # Get images
    target_imgs = sorted([f for f in INPUT_DIR.iterdir() if f.suffix.lower() in [".jpg", ".png", ".jpeg"]])
    
    print(f"Found {len(target_imgs)} target images.")
    
    for img_path in target_imgs:
        for style in STYLES:
            run_job(img_path, style)
            time.sleep(1) # Small pause between submissions

if __name__ == "__main__":
    main()