import io
import os
import modal

# ==============================================================================
# 💾 0. CACHE VOLUME SETUP
# ==============================================================================
# This creates a persistent volume that survives between container cold-starts.
cache_volume = modal.Volume.from_name("flux-inductor-cache", create_if_missing=True)
CACHE_DIR = "/root/.cache/torch/inductor"

# Define container environment
image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "diffusers",
        "transformers",
        "accelerate",
        "pillow",
        "torch",
        "triton"  # Essential for torch.compile
    )
    # Tell PyTorch Inductor to write its cache to our persistent volume directory
    .env({"TORCHINDUCTOR_CACHE_DIR": CACHE_DIR})
)

app = modal.App("flux-klein-voxel-backend", image=image)

# ==============================================================================
# 🏎️ 1. THE DEMO PIPELINE (FALLBACK ROUTE)
# ==============================================================================
@app.function()
def demo_stream_frame(img_bytes: bytes) -> bytes:
    """Fallback route structurally aligned to match the frontend signature."""
    from PIL import Image, ImageDraw
    
    input_image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    draw = ImageDraw.Draw(input_image)
    draw.text((20, 20), "🛠️ WEBRTC PASSTHROUGH DEMO ACTIVE", fill=(0, 255, 0))
    
    output_buffer = io.BytesIO()
    input_image.save(output_buffer, format="JPEG", quality=85)
    return output_buffer.getvalue()


# ==============================================================================
# 🚀 2. THE REAL-TIME VOXEL ENGINE
# ==============================================================================
@app.cls(
    gpu="A10G",
    secrets=[modal.Secret.from_name("huggingface-secret")],
    max_containers=5,
    # Mount the persistent cache volume to the exact path PyTorch is looking at
    volumes={CACHE_DIR: cache_volume} 
)
class VoxelModel:
    
    @modal.enter()
    def load_pipeline(self):
        import torch
        from PIL import Image
        from diffusers import AutoPipelineForImage2Image
        
        model_id = "AnimeOverlord/flux2-klein-4b-mc"
        
        self.pipe = AutoPipelineForImage2Image.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
            token=os.environ["HF_TOKEN"]  # Modernized from use_auth_token
        )
        self.pipe.to("cuda")
        self.pipe.enable_attention_slicing()

        # ---------------------------------------------------------
        # ✨ COMPILE AND OPTIMIZE THE INFRASTRUCTURE
        # ---------------------------------------------------------
        print("Initializing torch.compile optimization loops...")
        
        # Channels-last optimization applied exclusively to the VAE (CNN-based)
        self.pipe.vae.to(memory_format=torch.channels_last)
        
        # Compile the heaviest part of the FLUX architecture safely without memory layout issues
        self.pipe.transformer = torch.compile(
            self.pipe.transformer,
            mode="reduce-overhead",  # Trades a bit of compile time for faster inference
            fullgraph=False
        )

        # WARMUP RUN: Force a dummy inference immediately execution starts.
        # If the cache is empty (first run), this traces the graph and saves to the Modal Volume.
        # If the cache exists, it instantly hotloads from the Volume.
        print("Running warmup to build/load inductor cache...")
        dummy_image = Image.new("RGB", (512, 512), (0, 0, 0))
        with torch.inference_mode():
            self.pipe(
                prompt="warmup pass",
                image=dummy_image,
                strength=0.5,
                num_inference_steps=2,
                guidance_scale=1.0,
            )
        print("Warmup complete. Ready for real-time requests!")

    @modal.method()
    def process_frame(self, img_bytes: bytes) -> bytes:
        from PIL import Image
        import torch
        
        # Hardcoded parameters that were previously passed from the UI
        prompt = "isometric 3d minecraft block voxel style, high resolution, volumetric lighting"
        strength = 0.45
        
        input_image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
        input_image = input_image.resize((512, 512)) 
        
        with torch.inference_mode():
            output_image = self.pipe(
                prompt=prompt,
                image=input_image,
                strength=strength,
                num_inference_steps=4,
                guidance_scale=3.5,
            ).images[0]
            
        output_buffer = io.BytesIO()
        output_image.save(output_buffer, format="JPEG", quality=85)
        return output_buffer.getvalue()