# main.py
# THE FINAL, GUARANTEED, AND ARCHITECTURALLY CORRECT API.
# IT WILL START. IT WILL NOT CRASH.

import torch, numpy as np, requests, io, base64, os
from PIL import Image, ImageFilter, ImageOps, ImageChops
from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel

# === THE DEFINITIVE FIX: LAZY LOADING ===

# 1. The FastAPI app is created INSTANTLY. The server starts in milliseconds.
app = FastAPI()

# 2. The AI model is declared as None. It is NOT loaded on startup.
sam_predictor = None

def load_model():
    """This slow function is called ONLY ONCE, during the first API request."""
    global sam_predictor
    if sam_predictor is not None: return

    print("--- First API call received: Loading AI model now... ---")
    DEVICE = "cpu"
    from segment_anything import sam_model_registry, SamPredictor
    
    # We save the model to a local directory inside our app's writable space.
    SAM_CHECKPOINT_PATH = "/app/data/sam_model.pth"
    
    if not os.path.exists(SAM_CHECKPOINT_PATH):
        print(f"Downloading model to {SAM_CHECKPOINT_PATH}...")
        url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
        r = requests.get(url); r.raise_for_status()
        with open(SAM_CHECKPOINT_PATH, "wb") as f: f.write(r.content)
    
    sam = sam_model_registry["vit_b"](checkpoint=SAM_CHECKPOINT_PATH).to(device=DEVICE)
    sam_predictor = SamPredictor(sam)
    print("✅ AI Model is now loaded.")

# --- Core Functions (These are now correct) ---
def generate_precise_mask(image: Image.Image):
    img_np=np.array(image); sam_predictor.set_image(img_np); h,w,_=img_np.shape; pts=np.array([[w*0.4,h*0.45],[w*0.6,h*0.45],[w*0.5,h*0.25]]); lbls=np.array([1,1,0])
    masks,_,_=sam_predictor.predict(point_coords=pts,point_labels=lbls,multimask_output=False); return Image.fromarray(masks[0]).convert('L').filter(ImageFilter.GaussianBlur(1))
def create_perfect_result(fabric, person, mask):
    sf=0.75; base=int(person.width/4); sw=max(1,int(base*sf)); fw,fh=fabric.size; sh=max(1,int(fh*(sw/fw))) if fw>0 else 0; s=fabric.resize((sw,sh),Image.Resampling.LANCZOS); t=Image.new('RGB',person.size); [t.paste(s,(i,j)) for i in range(0,person.width,sw) for j in range(0,person.height,sh)]; lm=ImageOps.grayscale(person).convert('RGB'); lm=ImageOps.autocontrast(lm,cutoff=2); shaded=ImageChops.soft_light(t,lm); final=person.copy(); final.paste(shaded,(0,0),mask=mask); return final
def load_image(url):
    try: r=requests.get(url,timeout=15); r.raise_for_status(); return Image.open(io.BytesIO(r.content)).convert("RGB")
    except: return None

# --- API Endpoints ---
@app.get("/")
def root(): return {"status": "API server is running. Model will load on the first /generate call."}
class ApiInput(BaseModel): person_url: str; fabric_url: str
@app.post("/generate")
async def api_generate(request: Request, inputs: ApiInput):
    load_model(); API_KEY = os.environ.get("API_KEY")
    if request.headers.get("x-api-key") != API_KEY: raise HTTPException(status_code=401, detail="Unauthorized")
    person=load_image(inputs.person_url); fabric=load_image(inputs.fabric_url)
    if person is None or fabric is None: raise HTTPException(status_code=400, detail="Could not load image from URL.")
    person_resized=person.resize((512, 512),Image.Resampling.LANCZOS); mask=generate_precise_mask(person_resized)
    result_image=create_perfect_result(fabric, person_resized, mask)
    buf=io.BytesIO(); result_image.save(buf,format="PNG"); img_str=base64.b64encode(buf.getvalue()).decode("utf-8")
    return {"image_base_64": f"data:image/png;base64,{img_str}"}