import torch
import logging
import io
from pathlib import Path
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from starlette.responses import StreamingResponse
from diffusers import StableDiffusionPipeline
from peft import PeftModel
from fastapi.middleware.cors import CORSMiddleware





# ==================================
# Setup: App, Logging, Constants
# ==================================
app = FastAPI(
    title="Stable Diffusion LoRA Inference API",
    description="The Flake API to generate images using a Stable Diffusion model fine-tuned with LoRA.",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Or specify your frontend's URL
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)

# Constants for model paths
BASE_MODEL_ID = "runwayml/stable-diffusion-v1-5"
LORA_UNET_PATH = "./lora_unet"
LORA_TEXT_ENCODER_PATH = "./lora_text"

# Global variable to hold the pipeline
pipe = None

# ==================================
# Pydantic Models for API
# ==================================
class ImageRequest(BaseModel):
    """Request model for image generation."""
    prompt: str
    num_steps: int = 30
    guidance_scale: float = 7.5

# ==================================
# Model Loading & Application Events
# ==================================
@app.on_event("startup")
def load_model():
    """
    Load the model and adapters once when the application starts.
    This is efficient as the model stays in memory.
    """
    global pipe
    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch_dtype = torch.float16 if device == "cuda" else torch.float32
    logging.info(f"Using device: {device}")

    # Load the base pipeline
    logging.info(f"Loading base model: {BASE_MODEL_ID}")
    pipe = StableDiffusionPipeline.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=torch_dtype,
        safety_checker=None,
    )
    
    # Load and fuse the LoRA adapters
    try:
        logging.info(f"Loading and attaching UNet LoRA from: {LORA_UNET_PATH}")
        pipe.unet = PeftModel.from_pretrained(pipe.unet, LORA_UNET_PATH)

        logging.info(f"Loading and attaching Text Encoder LoRA from: {LORA_TEXT_ENCODER_PATH}")
        pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, LORA_TEXT_ENCODER_PATH)
    except Exception as e:
        logging.error(f"Fatal: Could not load LoRA adapters. Server will not work. Error: {e}")
        # In a real scenario, you might want the app to fail startup if models don't load.
        raise RuntimeError(f"Failed to load LoRA adapters: {e}") from e

    pipe.to(device)
    logging.info("Model and LoRA adapters loaded successfully.")


# ==================================
# API Endpoints
# ==================================
@app.get("/")
def read_root():
    """Root endpoint to check if the API is running."""
    return {"status": "API is running", "docs_url": "/docs"}

@app.post("/generate-image")
async def generate_image_endpoint(request: ImageRequest):
    """
    Endpoint to generate an image.
    Receives a POST request with a JSON body containing the prompt.
    Returns the generated image as a PNG file.
    """
    if pipe is None:
        raise HTTPException(status_code=503, detail="Model is not loaded or failed to load.")

    try:
        logging.info(f"Generating image for prompt: '{request.prompt}'")
        image = pipe(
            prompt=request.prompt,
            num_inference_steps=request.num_steps,
            guidance_scale=request.guidance_scale,
        ).images[0]
        
        # Save image to an in-memory buffer
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        buffer.seek(0)

        # Return the image as a streaming response
        return StreamingResponse(buffer, media_type="image/png")

    except Exception as e:
        logging.error(f"Error during image generation: {e}")
        raise HTTPException(status_code=500, detail=str(e))