# ----------------------------------------------------------------------
# IMPORTS
# ----------------------------------------------------------------------
import torch
import logging
import time
import spaces
import sys
import traceback
from PIL import Image
from typing import List, Optional, Any
from collections import defaultdict
from src.utils import LOG_LEVEL_MAP, EMOJI_MAP

# ----------------------------------------------------------------------
# RT-DETR CONSTANTS
# ----------------------------------------------------------------------
RTDETR_CONF = 0.4
RTDETR_ARTIFACT_CONF = 0.35

# ----------------------------------------------------------------------
# MODEL LABEL CONFIGURATION
# ----------------------------------------------------------------------
MODEL_LABEL_CONFIG = {
    "rtdetr_model": {
        "person_list": {
            "person": ["person"]
        },
        "product_type_list": {},
        "head_list": {},
        "shoes_list": {},
        "clothing_features_list": {
            "collar": ["tie"]
        },
        "artifacts_list": {
            "bag": ["backpack", "handbag", "suitcase"],
            "cup": ["bottle", "wine glass", "cup"],
            "umbrella": ["umbrella"],
            "book": ["book"],
            "phone": ["cell phone"],
            "camera": [],
            "other": ["fork", "knife", "spoon", "bowl", "frisbee", "sports ball", "kite",
                     "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
                     "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv",
                     "laptop", "mouse", "remote", "keyboard", "microwave", "oven", "toaster",
                     "sink", "refrigerator", "clock", "vase", "scissors", "teddy bear",
                     "hair drier", "toothbrush"]
        }
    }
}

# ----------------------------------------------------------------------
# RT-DETR HELPER FUNCTIONS
# ----------------------------------------------------------------------
def get_rtdetr_clothing_labels():
    clothing_labels = set()
    rtdetr_config = MODEL_LABEL_CONFIG.get("rtdetr_model", {})
    
    for keyword, labels in rtdetr_config.get("person_list", {}).items():
        clothing_labels.update(labels)
    
    for keyword, labels in rtdetr_config.get("product_type_list", {}).items():
        clothing_labels.update(labels)
    
    clothing_labels.update(["coat", "dress", "jacket", "shirt", "skirt", "pants", "shorts"])
    
    return clothing_labels

def get_rtdetr_person_and_product_labels():
    labels = set()
    rtdetr_config = MODEL_LABEL_CONFIG.get("rtdetr_model", {})
    
    for keyword, label_list in rtdetr_config.get("person_list", {}).items():
        labels.update(label_list)
    
    for keyword, label_list in rtdetr_config.get("product_type_list", {}).items():
        labels.update(label_list)
    
    labels.update(["person", "coat", "dress", "jacket", "shirt", "skirt", "pants", "shorts"])
    
    return labels

def get_rtdetr_artifact_labels():
    artifact_labels = set()
    rtdetr_config = MODEL_LABEL_CONFIG.get("rtdetr_model", {})
    
    for keyword, labels in rtdetr_config.get("artifacts_list", {}).items():
        if keyword != "other":
            artifact_labels.update(labels)
    
    return artifact_labels

def get_label_name_from_model(model, label_id):
    if hasattr(model, 'config') and hasattr(model.config, 'id2label'):
        return model.config.id2label.get(label_id, f"unknown_{label_id}").lower()
    if hasattr(model, 'model_labels') and isinstance(model.model_labels, dict):
        return model.model_labels.get(label_id, f"unknown_{label_id}").lower()
    return f"unknown_{label_id}"

def map_label_to_keyword(label_name: str, valid_kws: List[str], model_name: str) -> Optional[str]:
    ln = label_name.strip().lower()
    
    model_config = MODEL_LABEL_CONFIG.get(model_name, {})
    
    for list_type in ["person_list", "product_type_list", "head_list", 
                      "shoes_list", "clothing_features_list", "artifacts_list"]:
        category_config = model_config.get(list_type, {})
        
        for keyword, labels in category_config.items():
            if keyword in valid_kws:
                for label in labels:
                    if ln == label.lower() or ln in label.lower():
                        return keyword
    
    return None

def process_rtdetr_results(results, model, label_set, threshold, fallback_box=None):
    try:
        if isinstance(results, list):
            if len(results) > 0:
                result = results[0]
            else:
                return None, 0.0, None
        else:
            result = results
        
        found_box = None
        found_score = 0.0
        found_label = None
        
        for score, label, box in zip(result["scores"], result["labels"], result["boxes"]):
            score_val = score.item()
            if score_val < threshold:
                continue
                
            label_id = label.item()
            label_name = get_label_name_from_model(model, label_id)
            
            if label_name in label_set:
                x1, y1, x2, y2 = [int(val) for val in box.tolist()]
                found_box = [x1, y1, x2, y2]
                found_score = score_val
                found_label = label_name
                break
        
        return found_box, found_score, found_label
    except Exception as e:
        logging.log(LOG_LEVEL_MAP["WARNING"], f"{EMOJI_MAP['WARNING']} Error processing RTDETR results: {e}")
        return fallback_box, 0.0, None

# ----------------------------------------------------------------------
# RT-DETR DETECTION FUNCTIONS
# ----------------------------------------------------------------------
def detect_rtdetr_in_roi(roi_rgb, RTDETR_PROCESSOR, RTDETR_MODEL, DEVICE, log_item):
    boxes = []
    labels = []
    scores = []
    raw_labels = []
    
    try:
        rtdetr_inputs = RTDETR_PROCESSOR(images=roi_rgb, return_tensors="pt")
        rtdetr_inputs = {k: v.to(DEVICE) for k, v in rtdetr_inputs.items()}
        
        with torch.no_grad():
            rtdetr_outputs = RTDETR_MODEL(**rtdetr_inputs)
        
        rtdetr_results = RTDETR_PROCESSOR.post_process_object_detection(
            rtdetr_outputs,
            target_sizes=torch.tensor([[roi_rgb.height, roi_rgb.width]]).to(DEVICE),
            threshold=RTDETR_CONF
        )
        
        if isinstance(rtdetr_results, list) and len(rtdetr_results) > 0:
            result = rtdetr_results[0]
        else:
            result = rtdetr_results
            
        for score, label, box in zip(result["scores"], result["labels"], result["boxes"]):
            label_id = label.item()
            score_val = score.item()
            x1, y1, x2, y2 = [int(val) for val in box.tolist()]
            label_name = get_label_name_from_model(RTDETR_MODEL, label_id)
            
            boxes.append([x1, y1, x2, y2])
            labels.append(label_id)
            scores.append(score_val)
            raw_labels.append(label_name)
            
            logging.log(LOG_LEVEL_MAP["INFO"], f"rtdetr_model: {EMOJI_MAP['INFO']} RT-DETR detected: {label_name} at score {score_val:.3f}")
            
    except Exception as e:
        error_msg = f"RTDETR detection error: {str(e)}"
        error_trace = traceback.format_exc()
        
        logging.log(LOG_LEVEL_MAP["WARNING"], f"{EMOJI_MAP['WARNING']} {error_msg}")
        logging.error(f"Traceback:\n{error_trace}")
        
        log_item["warnings"] = log_item.get("warnings", []) + [error_msg]
        log_item["traceback"] = error_trace
        
        if "CUDA must not be initialized" in str(e):
            logging.critical("CUDA initialization error in Spaces Zero GPU environment")
            sys.exit(1)
    
    return boxes, labels, scores, raw_labels

def detect_rtdetr_artifacts_in_roi(roi_rgb, keywords, RTDETR_PROCESSOR, RTDETR_MODEL, DEVICE, log_item):
    boxes = []
    labels = []
    scores = []
    raw_labels = []
    
    try:
        rtdetr_inputs = RTDETR_PROCESSOR(images=roi_rgb, return_tensors="pt")
        rtdetr_inputs = {k: v.to(DEVICE) for k, v in rtdetr_inputs.items()}
        
        with torch.no_grad():
            rtdetr_outputs = RTDETR_MODEL(**rtdetr_inputs)
        
        rtdetr_results = RTDETR_PROCESSOR.post_process_object_detection(
            rtdetr_outputs,
            target_sizes=torch.tensor([[roi_rgb.height, roi_rgb.width]]).to(DEVICE),
            threshold=RTDETR_ARTIFACT_CONF
        )
        
        rtdetr_artifact_labels = get_rtdetr_artifact_labels()
        
        if isinstance(rtdetr_results, list) and len(rtdetr_results) > 0:
            result = rtdetr_results[0]
        else:
            result = rtdetr_results
            
        for score, label, box in zip(result["scores"], result["labels"], result["boxes"]):
            label_id = label.item()
            score_val = score.item()
            
            if score_val < RTDETR_ARTIFACT_CONF:
                continue
                
            label_name = get_label_name_from_model(RTDETR_MODEL, label_id)
            
            if label_name in rtdetr_artifact_labels:
                x1, y1, x2, y2 = [int(val) for val in box.tolist()]
                
                artifact_keyword = map_label_to_keyword(label_name, keywords, "rtdetr_model")
                if not artifact_keyword:
                    continue
                    
                boxes.append([x1, y1, x2, y2])
                labels.append(label_id)
                scores.append(score_val)
                raw_labels.append(label_name)
                
                logging.log(LOG_LEVEL_MAP["INFO"], f"rtdetr_model: {EMOJI_MAP['INFO']} Artifact detected: {label_name} at score {score_val:.3f}")
                
    except Exception as e:
        error_msg = f"RTDETR artifact detection error: {str(e)}"
        error_trace = traceback.format_exc()
        
        logging.log(LOG_LEVEL_MAP["WARNING"], f"{EMOJI_MAP['WARNING']} {error_msg}")
        logging.error(f"Traceback:\n{error_trace}")
        
        log_item["warnings"] = log_item.get("warnings", []) + [error_msg]
        log_item["traceback"] = error_trace
        
        if "CUDA must not be initialized" in str(e):
            logging.critical("CUDA initialization error in Spaces Zero GPU environment")
            sys.exit(1)
    
    return boxes, labels, scores, raw_labels

def update_fallback_detection(ctx, pi_rgba, fallback_box, RTDETR_PROCESSOR, RTDETR_MODEL, DEVICE, RTDETR_CONF, final_boxes, final_labels, final_scores, final_kws, final_raws, final_mods, dd_log):
    try:
        if not (fallback_box and isinstance(fallback_box, list) and len(fallback_box) == 4):
            return final_boxes, final_labels, final_scores, final_kws, final_raws, final_mods, dd_log
            
        sub_ = pi_rgba.crop((
            fallback_box[0],
            fallback_box[1],
            fallback_box[2],
            fallback_box[3]
        ))
        sub_ = sub_.convert("RGB")
        subW = sub_.width
        subH = sub_.height
        
        rtdetr_inputs = RTDETR_PROCESSOR(images=sub_, return_tensors="pt").to(DEVICE)
        
        with torch.no_grad():
            rtdetr_outputs = RTDETR_MODEL(**rtdetr_inputs)
        
        rtdetr_results = RTDETR_PROCESSOR.post_process_object_detection(
            rtdetr_outputs,
            target_sizes=torch.tensor([[subH, subW]]).to(DEVICE),
            threshold=RTDETR_CONF
        )
        
        rtdetr_clothing_labels = get_rtdetr_clothing_labels()
        found_fb_box, found_fb_score, _ = process_rtdetr_results(
            rtdetr_results, RTDETR_MODEL, rtdetr_clothing_labels, RTDETR_CONF
        )
        
        if found_fb_box:
            fx1 = fallback_box[0] + found_fb_box[0]
            fy1 = fallback_box[1] + found_fb_box[1]
            fx2 = fallback_box[0] + found_fb_box[2]
            fy2 = fallback_box[1] + found_fb_box[3]
            
            found_fb_box = [fx1, fy1, fx2, fy2]
            final_boxes.append(found_fb_box)
            final_labels.append(90001)
            final_scores.append(round(found_fb_score, 2))
            final_kws.append(ctx.product_type)
            final_raws.append("fallback_label")
            final_mods.append("rtdetr_model")
            dd_log[ctx.product_type].append({
                "box": found_fb_box,
                "score": found_fb_score,
                "raw_label": "fallback_label",
                "model": "rtdetr_model"
            })
        else:
            final_boxes.append(fallback_box)
            final_labels.append(90000)
            final_scores.append(0.0)
            final_kws.append(ctx.product_type)
            final_raws.append("fallback_label")
            final_mods.append("fallback")
            dd_log[ctx.product_type].append({
                "box": fallback_box,
                "score": 0.0,
                "raw_label": "fallback_label",
                "model": "fallback"
            })
        
        return final_boxes, final_labels, final_scores, final_kws, final_raws, final_mods, dd_log
    except Exception as e:
        logging.log(LOG_LEVEL_MAP["WARNING"], f"{EMOJI_MAP['WARNING']} Fallback detection error: {e}")
        final_boxes.append(fallback_box)
        final_labels.append(90000)
        final_scores.append(0.0)
        final_kws.append(ctx.product_type)
        final_raws.append("fallback_label")
        final_mods.append("fallback_error")
        dd_log[ctx.product_type].append({
            "box": fallback_box,
            "score": 0.0,
            "raw_label": "fallback_label",
            "model": "fallback_error"
        })
        return final_boxes, final_labels, final_scores, final_kws, final_raws, final_mods, dd_log