from fastapi import FastAPI
from pydantic import BaseModel
import torch
from diffusers import StableDiffusionPipeline
import story_generator  # Import Story Generator

app = FastAPI()

# ✅ Use a single request format
class StoryRequest(BaseModel):
    theme: str
    reading_level: str

# ✅ Load Image Model (Cartoon-Style)
CACHE_DIR = "/tmp/huggingface"
pipe = StableDiffusionPipeline.from_pretrained(
    "nitrosocke/Arcane-Diffusion",
    cache_dir=CACHE_DIR, 
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)
pipe.to("cuda" if torch.cuda.is_available() else "cpu")

@app.post("/generate_story_and_image")
def generate_story_and_image(request: StoryRequest):
    """Generates a story and an image based on the given theme and reading level."""
    print(f"Generating story for theme: {request.theme} and level: {request.reading_level}")

    # ✅ Generate story
    story_result = story_generator.generate_story_and_questions(request.theme, request.reading_level)
    story_text = story_result["story"]
    
    # ✅ Generate image
    cartoon_prompt = f"A colorful, cartoon-style illustration of: {story_text}, vibrant colors, highly detailed, storybook fantasy."
    print("Generating Image for:", cartoon_prompt[:100])
    image = pipe(cartoon_prompt, width=1024, height=1024).images[0]

    # ✅ Save image
    image_path = "generated_image.png"
    image.save(image_path)
    print("Image Saved!")  

    return {
        "theme": request.theme,
        "reading_level": request.reading_level,
        "story": story_text,
        "questions": story_result["questions"],
        "image_path": image_path
    }

# ✅ Fix the "Not Found" issue for the root URL
@app.get("/")
def home():
    return {"message": "Welcome to the Story & Image Generation API! Use /generate_story_and_image"}