import os
import random
import tempfile
import base64
import traceback
from io import BytesIO
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from PIL import Image, UnidentifiedImageError
from gradio_client import Client, handle_file
from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch
import logging
import time

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

# ========== CONFIG ==========
HF_API_NAME = "Seniordev22/face-swap-app2k6"
TARGET_SIZE = (768, 768)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MAX_RETRIES = 2

# ========== GLOBALS ==========
client = None
emotion_targets = {}          # key: (gender, emotion), value: list of image paths
feature_extractor = None
emotion_model = None
EMOTION_LABELS = []           # will be filled from model's id2label

# ========== MODEL LOADING (with dynamic label extraction) ==========
def load_emotion_model():
    global feature_extractor, emotion_model, EMOTION_LABELS
    if emotion_model is not None:
        return
    logger.info(f"Loading emotion model on {DEVICE}...")
    try:
        model_name = "dima806/facial_emotions_image_detection"
        feature_extractor = AutoImageProcessor.from_pretrained(model_name)
        emotion_model = AutoModelForImageClassification.from_pretrained(model_name)
        emotion_model.eval()
        emotion_model.to(DEVICE)

        # Extract correct label mapping from model config
        id2label = emotion_model.config.id2label
        # id2label is a dict like {0: 'angry', 1: 'disgust', ...}
        EMOTION_LABELS = [id2label[i] for i in range(len(id2label))]
        logger.info(f"✅ Emotion model loaded. Labels: {EMOTION_LABELS}")
    except Exception as e:
        logger.error(f"Failed to load emotion model: {e}")
        raise

def load_all_targets():
    """Load target images organized as: target/{male,female}/{emotion}/image.jpg"""
    global emotion_targets
    emotion_targets.clear()
    base_dir = "target"
    if not os.path.exists(base_dir):
        logger.warning("⚠️ 'target/' folder not found! Please create it with subfolders: male/female/emotion/images")
        return

    # Use the actual emotion labels from the model (not hardcoded)
    valid_emotions = set(EMOTION_LABELS)  # e.g., {'angry','disgust',...}
    loaded_count = 0

    for gender in ["male", "female"]:
        gender_path = os.path.join(base_dir, gender)
        if not os.path.exists(gender_path):
            logger.warning(f"⚠️ Missing folder: {gender_path}")
            continue

        for emotion in valid_emotions:
            emotion_dir = os.path.join(gender_path, emotion)
            if not os.path.exists(emotion_dir):
                logger.debug(f"No folder for {gender}/{emotion} – will fallback to random of same gender")
                continue

            images = []
            for f in os.listdir(emotion_dir):
                ext = os.path.splitext(f)[1].lower()
                if ext in {'.jpg', '.jpeg', '.png', '.webp', '.avif', '.bmp'}:
                    full_path = os.path.join(emotion_dir, f)
                    prepped = preprocess_image(full_path)
                    if prepped:
                        images.append(prepped)
            if images:
                emotion_targets[(gender, emotion)] = images
                loaded_count += len(images)
                logger.info(f"   Loaded {len(images)} images for {gender}/{emotion}")

    logger.info(f"🎯 Total target images loaded: {loaded_count}")
    if loaded_count == 0:
        logger.error("No target images found! Please populate target/male/emotion/ and target/female/emotion/")

# ========== HELPERS ==========
def preprocess_image(image_path: str):
    """Resize image if too large, return path (may be new temp path)"""
    try:
        img = Image.open(image_path).convert("RGB")
        if max(img.size) > max(TARGET_SIZE):
            img.thumbnail(TARGET_SIZE, Image.Resampling.LANCZOS)
            # Save to a temporary file to keep original unchanged
            fd, temp_path = tempfile.mkstemp(suffix=".jpg")
            os.close(fd)
            img.save(temp_path, "JPEG", quality=88, optimize=True)
            return temp_path
        return image_path
    except Exception as e:
        logger.warning(f"Preprocess failed for {image_path}: {e}")
        return None

def get_emotion(img_path: str) -> str:
    """Detect emotion from image file (no caching because temp files are unique)"""
    try:
        image = Image.open(img_path).convert("RGB")
        inputs = feature_extractor(images=image, return_tensors="pt")
        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = emotion_model(**inputs)
        pred_idx = outputs.logits.argmax(-1).item()
        emotion = EMOTION_LABELS[pred_idx]
        logger.debug(f"Detected emotion: {emotion} (idx {pred_idx})")
        return emotion
    except Exception as e:
        logger.error(f"Emotion detection failed: {e}")
        return "neutral"   # safe fallback

def safe_cleanup(path: str):
    try:
        if path and os.path.exists(path):
            os.unlink(path)
    except Exception:
        pass

def swap_face(source_img: Image.Image, gender: str):
    """Main face swap logic with emotion matching"""
    temp_files = []   # track all temp files for cleanup

    try:
        # Save source image to temp file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
            source_img = source_img.convert("RGB")
            source_img.save(tmp.name, "JPEG", quality=92)
            src_path = tmp.name
            temp_files.append(src_path)

        src_fixed = preprocess_image(src_path)
        if not src_fixed:
            return None, "Failed to preprocess source image"
        if src_fixed != src_path:
            temp_files.append(src_fixed)  # track if preprocess created a new temp

        # Detect emotion
        emotion = get_emotion(src_fixed)
        gender_key = "male" if gender.lower() in ["boy", "male"] else "female"
        logger.info(f"Source → Gender: {gender_key}, Emotion: {emotion}")

        # Pick target image based on (gender, emotion)
        target_paths = emotion_targets.get((gender_key, emotion))
        if not target_paths:
            # Fallback: any image of the same gender
            all_for_gender = []
            for (g, _), paths in emotion_targets.items():
                if g == gender_key:
                    all_for_gender.extend(paths)
            if not all_for_gender:
                return None, f"No target images available for gender '{gender_key}'"
            target_fixed = random.choice(all_for_gender)
            logger.warning(f"No {emotion} target for {gender_key} → using random from same gender")
        else:
            target_fixed = random.choice(target_paths)
            logger.info(f"Using target: {os.path.basename(target_fixed)} (emotion: {emotion})")

        # Call Gradio face swap API (with retries)
        result_img_path = None
        for attempt in range(MAX_RETRIES + 1):
            try:
                result = client.predict(
                    handle_file(src_fixed),
                    handle_file(target_fixed),
                    1, 1,
                    api_name="/lambda_2"
                )
                # Extract path from result
                if isinstance(result, (list, tuple)):
                    img_path = result[0]
                else:
                    img_path = result
                if isinstance(img_path, dict):
                    img_path = img_path.get("path") or img_path.get("url") or str(img_path)
                result_img_path = img_path
                break
            except Exception as e:
                if attempt == MAX_RETRIES:
                    raise
                logger.warning(f"Gradio attempt {attempt+1} failed: {e}. Retrying...")
                time.sleep(1)

        swapped_img = Image.open(result_img_path)
        return swapped_img, f"Success (emotion: {emotion})"

    except Exception as e:
        logger.error(f"Swap failed: {str(e)}\n{traceback.format_exc()}")
        return None, f"Face swap error: {str(e)}"
    finally:
        # Cleanup all temporary files
        for f in temp_files:
            safe_cleanup(f)

# ========== LIFESPAN ==========
@asynccontextmanager
async def lifespan(app: FastAPI):
    global client
    try:
        load_emotion_model()       # sets EMOTION_LABELS
        load_all_targets()         # uses EMOTION_LABELS to scan folders
        client = Client(HF_API_NAME, verbose=False)
        logger.info("✅ Connected to Gradio Space")
        logger.info("🚀 Server is fully ready!")
    except Exception as e:
        logger.error(f"Startup failed: {e}")
    yield

# ========== FASTAPI ==========
app = FastAPI(title="Gender Swap API", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

class SwapRequest(BaseModel):
    image: str   # base64 encoded
    gender: str  # "male" or "female"

class SwapResponse(BaseModel):
    image: str    # base64 encoded result
    emotion: str  # detected emotion from source
    message: str

@app.post("/swap/gender", response_model=SwapResponse)
async def swap_gender(request: SwapRequest):
    start_time = time.time()

    try:
        # Decode base64 input
        try:
            image_data = base64.b64decode(request.image)
            source_img = Image.open(BytesIO(image_data)).convert("RGB")
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid base64 image: {str(e)}")

        # Perform swap
        result_img, message = swap_face(source_img, request.gender)
        if result_img is None:
            raise HTTPException(status_code=500, detail=message)

        # Convert result to base64
        buffered = BytesIO()
        result_img.save(buffered, format="JPEG", quality=85, optimize=True)
        result_b64 = base64.b64encode(buffered.getvalue()).decode()

        # Get emotion from source (reuse the image, but swap_face already logged it)
        # For response, we can re-detect or just pass what swap_face used.
        # Simpler: detect again from source image (no big cost)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
            source_img.save(tmp.name, "JPEG", quality=90)
            emotion = get_emotion(tmp.name)
            safe_cleanup(tmp.name)

        logger.info(f"✅ Swap completed in {time.time()-start_time:.2f}s | Emotion: {emotion}")
        return SwapResponse(image=result_b64, emotion=emotion, message=message)

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Unexpected error: {traceback.format_exc()}")
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

@app.get("/health")
def health():
    return {
        "status": "ok",
        "targets_loaded": sum(len(v) for v in emotion_targets.values()),
        "emotion_labels": EMOTION_LABELS
    }

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