#!/usr/bin/env python3
"""
오일 유출 감지 추론 서버 (GPU)
시립대 starsafire ResNet101+DANet 모델 기반

실행: uvicorn oil_inference_server:app --host 0.0.0.0 --port 8090
모델 파일 필요: ./V7_SPECIAL.py, ./epoch_165.pth (같은 디렉토리)
"""

import os
import io
import base64
import logging
from collections import Counter
from typing import Optional

import cv2
import numpy as np
from PIL import Image
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# ── MMSegmentation (지연 임포트 — 서버 시작 시 로드) ─────────────────────────
model = None
DEVICE = os.getenv("INFERENCE_DEVICE", "cuda:0")

CLASSES = ("background", "black", "brown", "rainbow", "silver")
PALETTE = [
    [0, 0, 0],        # 0: background
    [0, 0, 204],      # 1: black oil (에멀전)
    [180, 180, 180],  # 2: brown oil (원유)
    [255, 255, 0],    # 3: rainbow oil (박막)
    [178, 102, 255],  # 4: silver oil (극박막)
]
THICKNESS_MM = {
    1: 1.0,      # black
    2: 0.1,      # brown
    3: 0.0003,   # rainbow
    4: 0.0001,   # silver
}

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("oil-inference")

# ── FastAPI App ──────────────────────────────────────────────────────────────

app = FastAPI(title="Oil Spill Inference Server", version="1.0.0")

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


class InferenceRequest(BaseModel):
    image: str  # base64 encoded JPEG/PNG


class OilRegionResult(BaseModel):
    classId: int
    className: str
    pixelCount: int
    percentage: float
    thicknessMm: float


class InferenceResponse(BaseModel):
    mask: str  # base64 encoded uint8 array (values 0-4)
    width: int
    height: int
    regions: list[OilRegionResult]


# ── Model Loading ────────────────────────────────────────────────────────────

def load_model():
    """모델을 로드한다. 서버 시작 시 1회 호출."""
    global model
    try:
        from mmseg.apis import init_segmentor

        script_dir = os.path.dirname(os.path.abspath(__file__))
        config_path = os.path.join(script_dir, "V7_SPECIAL.py")
        checkpoint_path = os.path.join(script_dir, "epoch_165.pth")

        if not os.path.exists(config_path):
            logger.error(f"Config not found: {config_path}")
            return False
        if not os.path.exists(checkpoint_path):
            logger.error(f"Checkpoint not found: {checkpoint_path}")
            return False

        logger.info(f"Loading model on {DEVICE}...")
        model = init_segmentor(config_path, checkpoint_path, device=DEVICE)
        model.PALETTE = PALETTE
        logger.info("Model loaded successfully")
        return True

    except Exception as e:
        logger.error(f"Model loading failed: {e}")
        return False


@app.on_event("startup")
async def startup():
    success = load_model()
    if not success:
        logger.warning("Model not loaded — inference will be unavailable")


# ── Endpoints ────────────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    return {
        "status": "ok" if model is not None else "model_not_loaded",
        "device": DEVICE,
        "classes": list(CLASSES),
    }


@app.post("/inference", response_model=InferenceResponse)
async def inference(req: InferenceRequest):
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        # 1. Base64 → numpy array
        img_bytes = base64.b64decode(req.image)
        img_pil = Image.open(io.BytesIO(img_bytes)).convert("RGB")
        img_np = np.array(img_pil)

        # 2. 임시 파일로 저장 (mmseg inference_segmentor는 파일 경로 필요)
        import tempfile
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
            tmp_path = tmp.name
            cv2.imwrite(tmp_path, cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR))

        # 3. 추론
        from mmseg.apis import inference_segmentor
        result = inference_segmentor(model, tmp_path)
        seg_map = result[0]  # (H, W) uint8, values 0-4

        # 임시 파일 삭제
        os.unlink(tmp_path)

        h, w = seg_map.shape
        total_pixels = h * w

        # 4. 클래스별 통계
        counter = Counter(seg_map.flatten().tolist())
        regions = []
        for class_id in range(1, 5):  # 1-4 (skip background)
            count = counter.get(class_id, 0)
            if count > 0:
                regions.append(OilRegionResult(
                    classId=class_id,
                    className=CLASSES[class_id],
                    pixelCount=count,
                    percentage=round(count / total_pixels * 100, 2),
                    thicknessMm=THICKNESS_MM[class_id],
                ))

        # 5. 마스크를 base64로 인코딩
        mask_bytes = seg_map.astype(np.uint8).tobytes()
        mask_b64 = base64.b64encode(mask_bytes).decode("ascii")

        return InferenceResponse(
            mask=mask_b64,
            width=w,
            height=h,
            regions=regions,
        )

    except Exception as e:
        logger.error(f"Inference error: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


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