from fastapi import FastAPI, File, UploadFile, Request
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
import os, json, base64, re
from dotenv import load_dotenv
import requests

os.environ["HF_HOME"] = "./cache"
os.makedirs("./cache", exist_ok=True)
os.makedirs("images", exist_ok=True)

import torch
from diffusers import StableDiffusionPipeline
from PIL import Image

load_dotenv()
api_key = os.getenv("GOOGLE_API")
app = FastAPI()

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

app.mount("/images", StaticFiles(directory="images"), name="images")
app.mount("/templates", StaticFiles(directory="templates"), name="templates")
os.makedirs("images", exist_ok=True)

device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16 if device == "cuda" else torch.float32)
pipe.to(device)

def clean_filename(text):
    return re.sub(r'[^\w\-_\. ]', '_', text.strip().lower().replace(" ", "_"))

def generate_image(food_name):
    prompt = f"Professional food photography of {food_name}, top-down view, realistic lighting"
    image = pipe(prompt).images[0]
    file_path = f"images/{clean_filename(food_name)}.png"
    image.save(file_path)
    return f"/{file_path}"

def extract_menu_from_image(image_bytes):
    base64_image = base64.b64encode(image_bytes).decode('utf-8')
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={api_key}"
    prompt = """
    Extract the menu items from this image and return ONLY a JSON array like:
    [
      {
        "food": "Dish Name",
        "description": "Short description or empty string",
        "price": 10,
        "category": "Category"
      }
    ]
    """
    payload = {
        "contents": [{"parts": [{"text": prompt}, {"inline_data": {"mime_type": "image/jpeg", "data": base64_image}}]}],
        "generationConfig": {"responseMimeType": "application/json"}
    }
    headers = {'Content-Type': 'application/json'}
    try:
        res = requests.post(url, headers=headers, json=payload)
        res.raise_for_status()
        text = res.json()['candidates'][0]['content']['parts'][0]['text']
        return json.loads(text)
    except Exception as e:
        print("Error extracting:", e)
        return []

@app.get("/", response_class=HTMLResponse)
async def serve_home():
    return FileResponse("templates/index.html")

@app.post("/upload-json")
async def upload_json(menu_image: UploadFile = File(...)):
    image_bytes = await menu_image.read()
    menu_items = extract_menu_from_image(image_bytes)

    for item in menu_items:
        img_path = generate_image(item["food"])
        item["img_path"] = img_path

    return JSONResponse(content=menu_items)