import os
import base64
import json
import requests
import torch
import numpy as np
import cv2
from PIL import Image, ImageFilter
from scipy.ndimage import binary_dilation
from omegaconf import OmegaConf

# -----------------------------
# 1) הגדרת המפתח API של Gemini כפרמטר
# -----------------------------

SYSTEM_INST = """\
You are given an image. You must return information about the main character in the image.
Do not write anything else beyond this!

**Guidelines for identifying a character in the image:**
1. **Male:**
   - Infant (0–2) → "baby boy"
   - Toddler (2–5) → "toddler boy"
   - Child (6–11) → "boy"
   - Teenager (12–17) → "teen boy"
   - Young adul (18–35) → "young man"
   - adul (36–59) → "man"
   - Elderly (60+) → "elderly man"

2. **Female:**
   - Infant (0–2) → "baby girl"
   - Toddler (2–5) → "toddler girl"
   - Child (6–11) → "girl"
   - Teenager (12–17) → "teen girl"
   - Young adul (18–35) → "young woman"
   - adul (36–59) → "woman"
   - Elderly (60+) → "elderly woman"

3. **Unclear identification:**
   - Ambiguous character → "unidentified"
   - Ambiguous infant/toddler → "baby" or "toddler"

4. **No character in the image:**
   - Respond: "no person"

5. **Multiple characters:**
   - Identify the most central or prominent character.

Notes:
- If data is insufficient to classify → "insufficient data".
"""

conversation = []  # נשמור כאן את השיחה הנוכחית

female_keywords = {
    "baby girl", "toddler girl", "girl",
    "teen girl", "young woman", "woman",
    "elderly woman"
}

def is_female_from_text(gemini_text: str) -> bool:
    """בודק האם התשובה מ-Gemini מצביעה על אישה לפי מילות המפתח שהוגדרו."""
    return gemini_text.lower().strip() in female_keywords


def encode_image_to_base64(image: Image.Image) -> str:
    import io
    buffer = io.BytesIO()
    image.save(buffer, format='JPEG')
    encoded_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
    return encoded_str

def add_user_text(message: str):
    conversation.append({
        "role": "user",
        "parts": [
            {"text": message}
        ]
    })


def add_user_image_from_pil(image: Image.Image, mime_type: str = "image/jpeg"):
    encoded_str = encode_image_to_base64(image)
    conversation.append({
        "role": "user",
        "parts": [
            {
                "inline_data": {
                    "mime_type": mime_type,
                    "data": encoded_str
                }
            }
        ]
    })


def send_and_receive(api_key: str) -> str:
    url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent"
    params = {"key": api_key}
    headers = {"Content-Type": "application/json"}

    payload = {
        "systemInstruction": {
            "role": "system",
            "parts": [
                {"text": SYSTEM_INST}
            ]
        },
        "contents": conversation
    }

    response = requests.post(url, params=params, headers=headers, json=payload)
    if response.status_code != 200:
        print(f"[Gemini] שגיאה בסטטוס קוד: {response.status_code}")
        return "NO_ANSWER"

    resp_json = response.json()
    candidates = resp_json.get("candidates", [])
    if not candidates:
        print("[Gemini] לא התקבלה תשובה.")
        return "NO_ANSWER"

    model_content = candidates[0].get("content", {})
    model_parts = model_content.get("parts", [])
    if not model_parts:
        print("[Gemini] לא נמצא תוכן בתשובת המודל.")
        return "NO_ANSWER"

    model_text = model_parts[0].get("text", "").strip()
    conversation.append({
        "role": "model",
        "parts": [
            {"text": model_text}
        ]
    })
    return model_text


# -----------------------------
# 3) טעינת מודל YOLO
# -----------------------------
from ultralytics import YOLO
YOLO_MODEL_PATH = '../../models/yolo11m.pt'

try:
    yolo_model = YOLO(YOLO_MODEL_PATH)
    yolo_model.to("cpu")
    print("[YOLO] מודל YOLO נטען בהצלחה.")
except Exception as e:
    print(f"[YOLO] לא מצליח לטעון את המודל בנתיב: {YOLO_MODEL_PATH}. שגיאה: {e}")
    yolo_model = None

TARGET_CLASS = "person"
CONF_THRESHOLD = 0.2

# -----------------------------
# 4) הכנה ל-SAM2
# -----------------------------
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

# נתיבים יחסיים ל-Space של Hugging Face
SAM2_CHECKPOINT = "checkpoints/sam2.1_hiera_tiny.pt"
MODEL_CFG = "configs/sam2.1/sam2.1_hiera_t.yaml"

sam2_predictor = None  # אתחול כ-None
device = "cuda" if torch.cuda.is_available() else "cpu"

def load_sam2_model():
    """טוען את מודל SAM2 באופן גלובלי."""
    global sam2_predictor
    try:
        # טעינת המודל
        sam2_model = build_sam2(MODEL_CFG, SAM2_CHECKPOINT, device=device)
        sam2_predictor = SAM2ImagePredictor(sam2_model)
        print("[SAM2] מודל SAM2 נטען בהצלחה.")

    except FileNotFoundError as e:
        print(f"[ERROR] קובץ SAM2 לא נמצא: {e}")
        print(f"  - ודא שקובץ המודל '{SAM2_CHECKPOINT}' וקובץ הקונפיג '{MODEL_CFG}' קיימים בנתיבים הנכונים בתוך ה-Space שלך.")

    except Exception as e:
        print(f"[ERROR] שגיאה כללית בטעינת SAM2: {e}")
        print(f"  - סוג השגיאה: {type(e).__name__}")
        print(f"  - הודעת השגיאה: {e}")
        import traceback
        print(f"  - Traceback:")
        traceback.print_exc()
        print(f"  - בדוק את התאימות בין גרסאות הספריות (torch, torchvision, sam2) ואת תקינות קובץ המודל.")

# טעינת המודל בעת טעינת המודול
load_sam2_model()


# -----------------------------
# 5) פונקציית טשטוש
# -----------------------------
def blur_regions_with_mask(
    image: Image.Image,
    mask: np.ndarray,
    blur_radius=20,
    pixel_size=20,
    expansion_pixels=1
):
    processed_image = image.copy()
    img_np = np.array(processed_image)

    structure = np.ones((expansion_pixels, expansion_pixels), dtype=bool)
    expanded_mask = binary_dilation(mask, structure=structure)

    blurred_whole = processed_image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
    blurred_whole_np = np.array(blurred_whole)

    ys, xs = np.where(expanded_mask)
    if len(xs) == 0 or len(ys) == 0:
        return processed_image

    x_min, x_max = xs.min(), xs.max()
    y_min, y_max = ys.min(), ys.max()

    region = blurred_whole_np[y_min:y_max, x_min:x_max]

    from PIL import Image as PILImage
    small = PILImage.fromarray(region).resize(
        ((x_max - x_min) // pixel_size, (y_max - y_min) // pixel_size),
        resample=Image.BILINEAR
    )
    pixelated = small.resize((x_max - x_min, y_max - y_min), PILImage.NEAREST)
    pixelated_np = np.array(pixelated)

    combined = img_np.copy()
    mask_region = expanded_mask[y_min:y_max, x_min:x_max]
    combined[y_min:y_max, x_min:x_max][mask_region] = pixelated_np[mask_region]

    return Image.fromarray(combined)


# -----------------------------
# 6) הפונקציה המרכזית
# -----------------------------
def process_image(
    pil_image: Image.Image,
    gemini_api_key: str,
    progress_callback=None
) -> Image.Image:
    if not gemini_api_key:
        raise ValueError("מפתח API של Gemini אינו מוזן.")
    """
    פונקציה המקבלת תמונת PIL, מפתח API של Gemini, ומחזירה את התמונה לאחר טשטוש נשים.
    """

    if progress_callback is None:
        # אם לא הועברה פונקציה לעדכון התקדמות, ניצור פונקציה ריקה
        def progress_callback(x, desc=""):
            pass

    conversation.clear()
    add_user_text("Processing a new image (backend)!")

    # 1) שלב YOLO
    progress_callback(0.0, "מתחיל זיהוי אנשים (YOLO)...")
    if yolo_model is None:
        print("[process_image] מודל YOLO לא נטען כראוי.")
        return pil_image

    np_image = np.array(pil_image)
    results = yolo_model.predict(np_image)
    bboxes_person = []

    for result in results:
        boxes = result.boxes
        for box in boxes:
            cls_name = yolo_model.names[int(box.cls)]
            conf = box.conf.item()
            if cls_name == TARGET_CLASS and conf >= CONF_THRESHOLD:
                x1, y1, x2, y2 = box.xyxy[0]
                bboxes_person.append([int(x1), int(y1), int(x2), int(y2)])

    progress_callback(0.1, f"נמצאו {len(bboxes_person)} בוקסי 'person' ב-YOLO")

    # 2) שלב Gemini (עבור כל בוקס בנפרד)
    women_boxes = []
    n_bboxes = len(bboxes_person) if bboxes_person else 1
    for i, bbox in enumerate(bboxes_person, start=1):
        fraction = 0.1 + (0.5 * i / n_bboxes)  # נניח חצי מההתקדמות מוקצה ל-Gemini
        progress_callback(fraction, f"[Gemini] בודק בוקס #{i} מתוך {len(bboxes_person)}")

        x1, y1, x2, y2 = bbox
        cropped = pil_image.crop((x1, y1, x2, y2))

        add_user_image_from_pil(cropped)
        add_user_text("---")

        gemini_text = send_and_receive(gemini_api_key)
        if is_female_from_text(gemini_text):
            women_boxes.append(bbox)

    # 3) שלב SAM2 (עבור בוקסים של נשים)
    if sam2_predictor is None:
        print("[process_image] SAM2 לא זמין/נטען. מחזירים תמונה ללא טשטוש.")
        raise ValueError("SAM2 model is not loaded.")

    progress_callback(0.6, f"מתחיל פילוח SAM2 על {len(women_boxes)} נשים...")
    sam2_predictor.set_image(np.array(pil_image))

    women_masks = []
    n_women = len(women_boxes) if women_boxes else 1
    for j, bbox in enumerate(women_boxes, start=1):
        fraction = 0.6 + (0.3 * j / n_women)  # עדכון עד 90%
        progress_callback(fraction, f"[SAM2] מפלח בוקס #{j} מתוך {len(women_boxes)}")

        box_np = np.array([bbox])
        masks, scores, _ = sam2_predictor.predict(
            point_coords=None,
            point_labels=None,
            box=box_np,
            multimask_output=False,
        )

        if masks.ndim == 4 and masks.shape[1] == 1:
            mask = masks.squeeze(1)[0].astype(bool)
        elif masks.ndim == 3:
            mask = masks[0].astype(bool)
        else:
            raise ValueError(f"[SAM2] צורת masks לא צפויה: {masks.shape}")

        women_masks.append((bbox, mask))

    # 4) שלב טשטוש
    progress_callback(0.9, "מתחיל טשטוש האזורים המזוהים (Blur + פיקסול)...")
    final_image = pil_image.copy()
    for (bbox, mask) in women_masks:
        final_image = blur_regions_with_mask(final_image, mask)

    progress_callback(1.0, "סיימנו! מחזירים את התוצאה הסופית.")

    # המרת התמונה ל-base64
    encoded_image = encode_image_to_base64(final_image)
    return encoded_image