import os
import shutil
from fastapi import FastAPI, UploadFile, File, HTTPException
from pydantic import BaseModel

# Import the master pipelines we just built
from ml_pipeline import process_url_pipeline, process_media_pipeline

app = FastAPI(title="Tag & Trail ML Brain")

# Pydantic model for incoming URL requests
class URLRequest(BaseModel):
    text: str

@app.post("/predict_url")
async def predict_url(req: URLRequest):
    """Receives a URL, runs ML, and extracts metadata if safe."""
    if not req.text:
        raise HTTPException(status_code=400, detail="No text/URL provided.")
    
    try:
        result = process_url_pipeline(req.text)
        
        # Ensure we return a "prediction" key so your helpers.py can read it perfectly
        if "class" in result:
            result["prediction"] = result["class"].lower()
            
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"URL Pipeline Error: {str(e)}")

@app.post("/predict_pdf")
async def predict_pdf(file: UploadFile = File(...)):
    """Receives a PDF/Media file, converts, runs ML, and extracts metadata if safe."""
    
    # Hugging Face Spaces allows writing to /tmp/ for ephemeral storage
    temp_dir = "/tmp/tag_and_trail_downloads"
    os.makedirs(temp_dir, exist_ok=True)
    
    temp_path = os.path.join(temp_dir, file.filename)
    
    try:
        # Save the incoming file from Twilio/helpers.py
        with open(temp_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
            
        # Run your master media pipeline
        result = process_media_pipeline(
            file_path=temp_path, 
            mime=file.content_type, 
            original_url_or_name=file.filename
        )
        
        return result
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Media Pipeline Error: {str(e)}")
        
    finally:
        # ALWAYS clean up to prevent memory/storage leaks in Hugging Face
        if os.path.exists(temp_path):
            os.remove(temp_path)

# Root endpoint just for quick health checks
@app.get("/")
def health_check():
    return {"status": "Tag & Trail ML API is running smoothly!"}