import io
import logging
import traceback
import numpy as np
from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from PIL import Image
from rembg import new_session, remove

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=['*'], allow_credentials=True, allow_methods=['*'], allow_headers=['*'])

SESSION = None
def get_session():
    global SESSION
    if SESSION is None:
        SESSION = new_session('birefnet-general-lite')
    return SESSION

@app.get('/')
def home():
    return {'status': 'ok', 'engine': 'BiRefNet-Lite-Turbo'}

@app.post('/api/remove-bg')
async def remove_bg(file: UploadFile = File(...)):
    try:
        input_data = await file.read()
        img = Image.open(io.BytesIO(input_data)).convert('RGBA')
        W, H = img.size
        
        # 1. FAST MASK PASS (to find subject area)
        # We use a very small image to find where the person is
        low_res = img.resize((256, 256), Image.NEAREST)
        mask_low = remove(low_res, session=get_session(), only_mask=True)
        
        # Find Bounding Box coordinates
        mask_np = np.array(mask_low)
        coords = np.argwhere(mask_np > 0)
        
        if coords.size > 0:
            y_min, x_min = coords.min(axis=0)
            y_max, x_max = coords.max(axis=0)
            
            # Scale back to original size + padding
            pad = 20
            l = max(0, int(x_min * W / 256) - pad)
            t = max(0, int(y_min * H / 256) - pad)
            r = min(W, int(x_max * W / 256) + pad)
            b = min(H, int(y_max * H / 256) + pad)
            
            # 2. CROP SUBJECT AREA
            subject_crop = img.crop((l, t, r, b))
            
            # 3. HIGH QUALITY PASS ON CROP ONLY (Much faster because it's small)
            # Resize crop if it's still too huge for CPU
            if max(subject_crop.size) > 1024:
                subject_crop.thumbnail((1024, 1024), Image.LANCZOS)
                
            cutout_crop = remove(subject_crop, session=get_session())
            
            # 4. PASTE BACK ON TRANSPARENT CANVAS
            final_img = Image.new('RGBA', (W, H), (0, 0, 0, 0))
            # Resize cutout back to the crop rectangle size if thumbnail was used
            cutout_crop = cutout_crop.resize((r-l, b-t), Image.LANCZOS)
            final_img.paste(cutout_crop, (l, t), cutout_crop)
        else:
            # Fallback if no subject found
            final_img = remove(img.resize((800, 800)), session=get_session()).resize((W, H))

        buf = io.BytesIO()
        final_img.save(buf, format='PNG')
        buf.seek(0)
        return StreamingResponse(buf, media_type='image/png')

    except Exception as e:
        return {'error': str(e)}

if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host='0.0.0.0', port=7860)