import os
import io
import json
import random
import time
import math
import sqlite3
import base64
import threading
from threading import Thread

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import torch
from huggingface_hub import hf_hub_download

from diffusers import Ideogram4Pipeline
from diffusers.quantizers.bitsandbytes.bnb_quantizer import BnB4BitDiffusersQuantizer

# System environmental controls matching your original template
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

DB_FILE = "generation_cache_queue.db"
MODEL_ID = "ideogram-ai/ideogram-4-nf4"
LM_HEAD_REPO = "multimodalart/qwen3-vl-8b-instruct-lm-head"
MAX_SEED = 2**31 - 1

# --- Runtime patch to preserve pristine diffusers integration ---
def _check_quantized_param_shape(self, param_name, current_param, loaded_param):
    n = math.prod(tuple(current_param.shape))
    inferred_shape = (n,) if "bias" in param_name else ((n + 1) // 2, 1)
    if tuple(loaded_param.shape) != tuple(inferred_shape):
        raise ValueError(f"Expected flattened shape of {param_name} to be {inferred_shape}, got {tuple(loaded_param.shape)}.")
    return True

BnB4BitDiffusersQuantizer.check_quantized_param_shape = _check_quantized_param_shape

# Presets mapping (Exact matching from your source configuration)
MODES = {
    "turbo": dict(num_inference_steps=12, guidance_schedule=(7.0,) * 11 + (3.0,) * 1, mu=0.5, std=1.75),
    "default": dict(num_inference_steps=20, guidance_schedule=(7.0,) * 18 + (3.0,) * 2, mu=0.0, std=1.75),
    "quality": dict(num_inference_steps=48, guidance_schedule=(7.0,) * 45 + (3.0,) * 3, mu=0.0, std=1.5),
}

# --- Shared Isolated SQLite Queue System Engine ---
def initialize_shared_database():
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS job_queue (
            id TEXT PRIMARY KEY,
            prompt TEXT NOT NULL,
            mode TEXT,
            width INTEGER,
            height INTEGER,
            status TEXT,
            image_data TEXT,
            created_at REAL
        )
    ''')
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_status_time ON job_queue(status, created_at)')
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_prompt_lookup ON job_queue(prompt)')
    conn.commit()
    conn.close()

initialize_shared_database()

# --- Pipeline Loading ---
print("[System] Loading Ideogram 4 pipeline directly into execution context...")
t_start = time.perf_counter()
pipe = Ideogram4Pipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
pipe.transformer.dequantize()
pipe.unconditional_transformer.dequantize()
pipe.to("cuda")
print(f"[System] Pipeline load complete: {time.perf_counter() - t_start:.1f}s")

# --- Background Worker Thread (The Concurrency Shock-Absorber) ---
def async_queue_processing_worker():
    print("[Worker] Async processing loop initialized and monitoring queue storage.")
    while True:
        try:
            conn = sqlite3.connect(DB_FILE)
            cursor = conn.cursor()
            
            # Fetch oldest queued task entry
            cursor.execute("""
                SELECT id, prompt, mode, width, height 
                FROM job_queue 
                WHERE status='queued' 
                ORDER BY created_at ASC LIMIT 1
            """)
            row = cursor.fetchone()
            
            if not row:
                conn.close()
                time.sleep(1.0) # Rest thread briefly when queue is cold
                continue
                
            task_id, prompt, mode, width, height = row
            
            # Lock entry row context to executing status
            cursor.execute("UPDATE job_queue SET status='processing' WHERE id=?", (task_id,))
            conn.commit()
            conn.close()
            
            print(f"[Worker] Running pipeline job allocation for Task: {task_id}")
            preset = MODES.get(mode, MODES["default"])
            seed = random.randint(0, MAX_SEED)
            
            # Execute Prompt Upsampling Graft
            expanded_prompt = pipe.upsample_prompt(
                prompt, height=height, width=width, lm_head_repo_id=LM_HEAD_REPO
            )[0]
            
            # Run Base Inference Loop 
            generator = torch.Generator(device="cuda").manual_seed(seed)
            output = pipe(prompt=expanded_prompt, width=width, height=height, generator=generator, **preset)
            image_obj = output.images[0]
            
            # Standard Binary Base64 Encoding Operation
            buffered_stream = io.BytesIO()
            image_obj.save(buffered_stream, format="JPEG", quality=85)
            base64_string = base64.b64encode(buffered_stream.getvalue()).decode("utf-8")
            data_uri_payload = f"data:image/jpeg;base64,{base64_string}"
            
            # Update Database status with Completed string resource
            conn = sqlite3.connect(DB_FILE)
            cursor = conn.cursor()
            cursor.execute("UPDATE job_queue SET status='completed', image_data=? WHERE id=?", (data_uri_payload, task_id))
            conn.commit()
            conn.close()
            print(f"[Worker] Task {task_id} marked complete. Assets delivered to storage.")
            
        except Exception as worker_error:
            print(f"[Critical Worker Error]: {str(worker_error)}")
            time.sleep(2.0)

# Start execution processing loop in isolated global background thread
threading.Thread(target=async_queue_processing_worker, daemon=True).start()

# --- High-Concurrency Web API Gateway Layout ---
app = FastAPI(title="Ideogram 4 Elastic Distribution Node")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"]
)

class RequestSchema(BaseModel):
    prompt: str
    mode: str = "default"
    width: int = 1024
    height: int = 1024

@app.post("/api/generate")
async def register_generation_job(payload: RequestSchema):
    # Enforce safe dimension factors matching baseline requirement constraints
    w = (payload.width // 64) * 64
    h = (payload.height // 64) * 64
    selected_mode = payload.mode.lower() if payload.mode.lower() in MODES else "default"
    
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    
    # 1. Automatic Global Deduplication / Performance Optimization Check
    cursor.execute("SELECT image_data FROM job_queue WHERE prompt=? AND status='completed' LIMIT 1", (payload.prompt,))
    cached_record = cursor.fetchone()
    
    if cached_record:
        conn.close()
        return {"status": "completed", "task_id": "cached_hit", "image": cached_record[0], "cached": True}
        
    # 2. Append new job entry parameters to the async structural file database
    new_task_uuid = str(random.randint(10000000, 99999999))
    cursor.execute(
        "INSERT INTO job_queue (id, prompt, mode, width, height, status, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
        (new_task_uuid, payload.prompt, selected_mode, w, h, "queued", time.time())
    )
    conn.commit()
    conn.close()
    
    return {"status": "queued", "task_id": new_task_uuid, "cached": False}

@app.get("/api/status/{task_id}")
async def inspect_job_status(task_id: str):
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute("SELECT status, image_data FROM job_queue WHERE id=?", (task_id,))
    record = cursor.fetchone()
    conn.close()
    
    if not record:
        raise HTTPException(status_code=404, detail="Requested generation task token variant not found.")
        
    current_status, alternative_image_payload = record
    return {"status": current_status, "image": alternative_image_payload}