import os
import cv2
import numpy as np
import pickle
import faiss
from insightface.app import FaceAnalysis
from fastapi import FastAPI, UploadFile, File, HTTPException
from io import BytesIO

# 1. FastAPI 애플리케이션 인스턴스 생성
app = FastAPI(
    title="InsightFace Face Recognition API",
    description="InsightFace (buffalo_l)를 사용한 얼굴 인식 및 FAISS 인덱스 검색 API"
)

# 모델 및 인덱스 전역 변수
model = None
index = None
labels = None

# 🚀 환경 설정: Docker 환경에서는 모델/인덱스 파일을 프로젝트 디렉토리에 배치합니다.
# 인덱스 파일은 /app/data 폴더에 있다고 가정합니다.
LOAD_DIR = "data"
FAISS_INDEX_FILE = os.path.join(LOAD_DIR, "faiss_index_v2.index")
LABELS_FILE = os.path.join(LOAD_DIR, "faiss_labels_v2.pkl")

# ✅ 1. 모델 및 인덱스 준비 (서버 시작 시 한 번만 실행)
@app.on_event("startup")
async def startup_event():
    global model, index, labels
    print("🚀 서버 시작: InsightFace 모델 및 FAISS 인덱스 로딩 중...")

    # InsightFace ArcFace 모델 준비
    try:
        # CPUExecutionProvider 사용 (GPU가 없는 환경/Docker에 적합)
        model = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
        model.prepare(ctx_id=0)
        print("✅ InsightFace model (buffalo_l) 로딩 완료.")
    except Exception as e:
        print(f"❌ InsightFace 모델 로딩 실패: {e}")
        raise HTTPException(status_code=500, detail=f"모델 로딩 실패: {e}")

    # 인덱스 & 라벨 로딩
    try:
        if not os.path.exists(FAISS_INDEX_FILE) or not os.path.exists(LABELS_FILE):
             print(f"❌ FAISS 파일이 없습니다. 경로 확인: {LOAD_DIR}")
             raise FileNotFoundError(f"필요한 FAISS 인덱스 파일 혹은 라벨 파일이 없습니다.")
        
        index = faiss.read_index(FAISS_INDEX_FILE)
        with open(LABELS_FILE, "rb") as f:
            labels = pickle.load(f)
        print(f"✅ FAISS 인덱스 로딩 완료. 총 {index.ntotal}개의 임베딩 로드.")
    except Exception as e:
        print(f"❌ FAISS 인덱스 로딩 실패: {e}")
        raise HTTPException(status_code=500, detail=f"FAISS 로딩 실패: {e}")
    
# 🚀 얼굴 임베딩 추출 함수 (원본 임베딩만 추출)
def get_face_embedding(img_np):
    """
    Numpy 배열 형태의 이미지에서 얼굴 임베딩을 추출합니다.
    """
    global model
    if model is None:
        raise HTTPException(status_code=500, detail="모델이 초기화되지 않았습니다.")

    # BGR을 RGB로 변환 (InsightFace 모델은 RGB를 선호)
    img = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
    faces = model.get(img)
    
    if faces:
        return faces[0].embedding
    else:
        return None
    


# --- API 엔드포인트 정의 ---
# 3. 루트 엔드포인트 (GET /)
@app.get("/")
def read_root():
    result={"success":True,"data":None,"msg":""}
    try:
        result["data"]="ok"
        return result
    except Exception as e:
        result["success"] = False
        result["msg"]=f"server error. {e!r}"
        return result
    
# 🚀 얼굴 예측 API 엔드포인트
@app.post("/predict_person/")
async def predict_person(
    image: UploadFile = File(..., description="분석할 얼굴 이미지 파일"),
    top_k: int = 1
):
    """
    업로드된 이미지에서 얼굴을 인식하고, FAISS 인덱스를 사용하여 가장 유사한 인물을 예측합니다.
    """
    global index, labels

    # 1. 이미지 파일 읽기
    content = await image.read()
    np_array = np.frombuffer(content, np.uint8)
    img_np = cv2.imdecode(np_array, cv2.IMREAD_COLOR)

    if img_np is None:
        raise HTTPException(status_code=400, detail="이미지 파일을 읽을 수 없습니다.")

    # 2. 얼굴 임베딩 추출
    embedding = get_face_embedding(img_np)
    
    if embedding is None:
        raise HTTPException(status_code=404, detail="이미지에서 얼굴을 인식할 수 없습니다.")

    # 3. 임베딩 정규화
    embedding = embedding.astype('float32')
    embedding /= np.linalg.norm(embedding)

    # 쿼리 형식에 맞게 [1, D] 형태로 변환
    query_vector = np.array([embedding])

    # 4. 유사도 검색
    # top_k는 최대 인덱스 크기(index.ntotal)를 초과할 수 없습니다.
    k = min(top_k, index.ntotal)
    scores, indices = index.search(query_vector, k)

    # 5. 결과 포맷팅
    results = []
    for idx, score in zip(indices[0], scores[0]):
        # labels[idx]는 인덱스에 저장된 임베딩 중 가장 유사한 임베딩의 라벨
        results.append({
            "rank": len(results) + 1,
            "person_id": labels[idx],
            "similarity_score": float(f"{score:.4f}")
        })
        
    return {"filename": image.filename, "predictions": results}


if __name__ == "__main__":
    # --reload 옵션을 추가하여 코드가 변경될 때마다 자동 재시작되게 설정합니다.
    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)