import os
import torch
from transformers import AutoModelForVision2Seq, AutoProcessor
from PIL import Image
import numpy as np
from fastapi import FastAPI, UploadFile, File, HTTPException
from typing import List
import io
import gradio as gr
from datetime import datetime
import json
import re
from exif import Image as ExifImage

# Initialize FastAPI app with increased upload limit (10MB)
app = FastAPI()

# Load SmolVLM-Instruct model and processor
model_id = "HuggingFaceTB/SmolVLM-Instruct"
processor = AutoProcessor.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"), padding=True)
model = AutoModelForVision2Seq.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"))

# Default harmful objects list
default_harmful_objects = ["knife", "gun", "weapon", "blood", "syringe", "bomb", "blade"]

# Case folder storage (in-memory for demo; persist to disk in production)
case_folders = {}

# Cancellation flag
cancel_flag = False

# 🔍 Core Features
# ----------------
# Resize image to max 1024x1024 while preserving aspect ratio
def resize_image(image: Image.Image, max_size: int = 1024) -> Image.Image:
    try:
        image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
        return image
    except Exception as e:
        raise ValueError(f"Image resizing failed: {str(e)}")

# Extract EXIF metadata, including GPS and timestamp
def extract_metadata(image_data: bytes) -> dict:
    try:
        exif_img = ExifImage(io.BytesIO(image_data))
        metadata = {
            "timestamp": exif_img.get("datetime_original", "N/A"),
            "gps": {
                "latitude": exif_img.get("gps_latitude", "N/A"),
                "longitude": exif_img.get("gps_longitude", "N/A")
            },
            "camera": f"{exif_img.get('make', 'N/A')} {exif_img.get('model', 'N/A')}"
        }
        return metadata
    except Exception as e:
        return {"error": f"Metadata extraction failed: {str(e)}"}

# Validate custom harmful objects input
def validate_custom_harmful(custom_harmful: str) -> List[str]:
    if not custom_harmful or not custom_harmful.strip():
        return []
    try:
        custom_objects = [obj.strip().lower() for obj in re.split(r'[,;]', custom_harmful) if obj.strip()]
        return [obj for obj in custom_objects if obj and all(c.isalnum() or c.isspace() for c in obj)]
    except Exception as e:
        raise ValueError(f"Invalid custom harmful objects: {str(e)}")

# 🚔 Investigation-Specific Features
# --------------------------------
# Calculate threat score based on harmful objects and weights
def calculate_threat_score(harmful_objects: List[dict], custom_weights: dict) -> float:
    score = 0.0
    for obj in harmful_objects or []:
        if obj.get("object") not in ["None", "Error", None]:
            weight = custom_weights.get(obj["object"], 1.0)
            score += (obj.get("confidence", 0) / 100) * weight
    return min(score, 100.0)  # Cap at 100

# 🧠 Smart Analysis Tools
# ----------------------
# Keyword-based search in results
def keyword_search(results: List[dict], keyword: str) -> List[dict]:
    if not keyword or not keyword.strip():
        return results
    keyword = keyword.lower().strip()
    filtered_results = []
    for result in results or []:
        match = False
        for key, value in result.items():
            if key in ["description", "signs", "faces", "scene_context", "clothing", "activity"] and isinstance(value, str) and keyword in value.lower():
                match = True
            elif key == "harmful_objects" and any(keyword in obj.get("object", "").lower() for obj in value or []):
                match = True
            elif key == "objects" and any(keyword in obj.get("object", "").lower() for obj in value or []):
                match = True
        if match:
            filtered_results.append(result)
    return filtered_results

# Manage case folders
def manage_case_folder(case_name: str, results: List[dict], action: str = "add") -> str:
    if not case_name or not case_name.strip():
        return "Error: Case folder name cannot be empty."
    if action == "add":
        if case_name not in case_folders:
            case_folders[case_name] = []
        case_folders[case_name].extend(results or [])
        return f"Added {len(results or [])} images to case folder '{case_name}'."
    elif action == "view":
        return json.dumps(case_folders.get(case_name, []), indent=2) or f"No data in case folder '{case_name}'."
    elif action == "clear":
        if case_name in case_folders:
            del case_folders[case_name]
            return f"Cleared case folder '{case_name}'."
        return f"Case folder '{case_name}' not found."
    return "Invalid action."

# 🛠️ User Options & Controls
# --------------------------
# Cancel analysis function
def cancel_analysis():
    global cancel_flag
    cancel_flag = True
    return "Analysis cancellation requested. Please wait for the current operation to stop."

# Format results as a detailed report (optimized)
def format_detailed_report(results, timestamp, findings, keyword: str = ""):
    global cancel_flag
    filtered_results = keyword_search(results, keyword)
    report_lines = [f"# Investigation Report\n**Generated on**: {timestamp}\n\n## Key Findings\n"]
    
    # Add key findings from precomputed dictionary
    if findings["harmful_found"]:
        report_lines.append("- Harmful objects detected in one or more images.\n")
    if findings["high_similarity"]:
        report_lines.append("- High similarity (>0.8) detected between images.\n")
    if findings["faces_detected"]:
        report_lines.append("- Faces detected with identifiable attributes.\n")
    if findings["threats_detected"]:
        report_lines.append("- High threat scores (>50) detected in one or more images.\n")
    if not any(findings.values()):
        report_lines.append("- No critical findings detected.\n")
    
    if cancel_flag:
        return "".join(report_lines) + "\n**Report Generation Cancelled**"
    
    report_lines.append("\n## Analysis Details\n")
    for result in filtered_results or []:
        if cancel_flag:
            return "".join(report_lines) + "\n**Report Generation Cancelled**"
        report_lines.append(f"### Image ID: {result.get('image_id', 'N/A')}\n")
        if "metadata" in result:
            meta = result["metadata"]
            report_lines.append(f"- **Metadata**: Timestamp: {meta.get('timestamp', 'N/A')}, GPS: {meta.get('gps', {}).get('latitude', 'N/A')}, {meta.get('gps', {}).get('longitude', 'N/A')}, Camera: {meta.get('camera', 'N/A')}\n")
        if "description" in result:
            report_lines.append(f"- **Description**: {result.get('description', 'N/A')}\n")
        if "signs" in result:
            report_lines.append(f"- **Signs**: {result.get('signs', 'N/A')}\n")
        if "harmful_objects" in result:
            harmful_str = ", ".join([f"{obj.get('object', 'N/A')} ({obj.get('confidence', 0)}%)" for obj in result.get('harmful_objects', [])]) or "None"
            report_lines.append(f"- **Harmful Objects**: {harmful_str}\n")
        if "similarity_to_image_1" in result:
            report_lines.append(f"- **Similarity**: {result.get('similarity_to_image_1', 0):.2f}\n")
        if "faces" in result:
            report_lines.append(f"- **Faces**: {result.get('faces', 'N/A')}\n")
        if "objects" in result:
            objects_str = ", ".join([f"{obj.get('object', 'N/A')} at {obj.get('bbox', 'N/A')}" for obj in result.get('objects', [])]) or "None"
            report_lines.append(f"- **Objects**: {objects_str}\n")
        if "clothing" in result:
            report_lines.append(f"- **Clothing/Colors**: {result.get('clothing', 'N/A')}\n")
        if "scene_context" in result:
            report_lines.append(f"- **Scene**: {result.get('scene_context', 'N/A')}\n")
        if "activity" in result:
            report_lines.append(f"- **Activity**: {result.get('activity', 'N/A')}\n")
        if "threat_score" in result:
            report_lines.append(f"- **Threat Score**: {result.get('threat_score', 0):.1f}/100\n")
        if "annotation" in result:
            report_lines.append(f"- **Annotation**: {result.get('annotation', 'N/A')}\n")
        if "flag" in result and result["flag"]:
            report_lines.append(f"- **Flagged**: Yes\n")
        if "comments" in result:
            report_lines.append(f"- **Investigator Comments**: {result.get('comments', 'N/A')}\n")
        if "error" in result:
            report_lines.append(f"- **Error**: {result.get('error', 'N/A')}\n")
        report_lines.append("\n")
    
    return "".join(report_lines)

@app.post("/predict")
async def predict(
    files: List[UploadFile] = File(...),
    description: bool = True,
    signs: bool = True,
    harmful: bool = True,
    similarity: bool = True,
    faces: bool = False,
    objects: bool = False,
    scene: bool = False,
    metadata: bool = False,
    activity: bool = False,
    clothing: bool = False,
    threat_score: bool = False,
    custom_harmful: str = "",
    custom_weights: str = "",
    description_level: str = "detailed"
):
    global cancel_flag
    cancel_flag = False  # Reset cancellation flag
    if not files or len(files) > 3:
        raise HTTPException(status_code=400, detail="Must upload 1–3 images.")
    
    # Process custom harmful objects and weights
    harmful_objects = default_harmful_objects.copy()
    custom_objects = validate_custom_harmful(custom_harmful)
    harmful_objects.extend(custom_objects)
    
    weights = {obj: 1.0 for obj in harmful_objects}
    if custom_weights.strip():
        try:
            for pair in custom_weights.split(","):
                obj, weight = pair.split(":")
                weights[obj.strip().lower()] = float(weight.strip())
        except:
            raise HTTPException(status_code=400, detail="Invalid custom weights format. Use 'object:weight,object:weight'.")
    
    results = []
    image_embeddings = []
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    for idx, file in enumerate(files, 1):
        if cancel_flag:
            return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
        try:
            # Check file size (limit to 10MB)
            image_data = await file.read()
            if len(image_data) > 10 * 1024 * 1024:
                raise ValueError("Image file size exceeds 10MB limit.")
            
            # Read and resize image
            image = Image.open(io.BytesIO(image_data)).convert("RGB")
            image = resize_image(image, max_size=1024)
            
            result = {"image_id": f"Image_{idx}", "timestamp": timestamp}
            
            # Core Features
            if metadata:
                if cancel_flag:
                    return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
                result["metadata"] = extract_metadata(image_data)
            
            if description:
                if cancel_flag:
                    return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
                prompt_desc = "<image> Provide a " + ("brief description of the image." if description_level == "basic" else "detailed description of the image, including objects, colors, people, and environmental context.")
                inputs = processor(text=[prompt_desc], images=[image], return_tensors="pt", padding=True)
                outputs = model.generate(**inputs, max_new_tokens=256 if description_level == "basic" else 512)
                result["description"] = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip() or "No description generated."
            
            if signs:
                if cancel_flag:
                    return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
                prompt_ocr = "<image> Extract all visible text in the image, such as road signs, license plates, or billboards, with exact wording."
                inputs_ocr = processor(text=[prompt_ocr], images=[image], return_tensors="pt", padding=True)
                ocr_outputs = model.generate(**inputs_ocr, max_new_tokens=256)
                result["signs"] = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip() or "None detected"
            
            if harmful:
                if cancel_flag:
                    return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
                prompt_detect = f"<image> Identify any harmful objects ({', '.join(harmful_objects)}) in the image. List them explicitly and estimate confidence (0-100%) for each detection based on clarity."
                inputs_detect = processor(text=[prompt_detect], images=[image], return_tensors="pt", padding=True)
                detect_outputs = model.generate(**inputs_detect, max_new_tokens=256)
                detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
                harmful_detected = []
                for obj in harmful_objects:
                    if obj in detected_objects:
                        confidence = 90 if obj in detected_objects.split() else 60
                        harmful_detected.append({"object": obj, "confidence": confidence})
                result["harmful_objects"] = harmful_detected if harmful_detected else [{"object": "None", "confidence": 0}]
            
            if faces:
                if cancel_flag:
                    return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
                prompt_faces = "<image> Detect faces and estimate attributes such as age range (e.g., child, adult, senior), gender (male, female, unknown), and emotional cues (e.g., neutral, angry, scared)."
                inputs_faces = processor(text=[prompt_faces], images=[image], return_tensors="pt", padding=True)
                faces_outputs = model.generate(**inputs_faces, max_new_tokens=256)
                result["faces"] = processor.decode(faces_outputs[0], skip_special_tokens=True).replace(prompt_faces, "").strip() or "No faces detected"
            
            if objects:
                if cancel_flag:
                    return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
                prompt_objects = "<image> Identify and localize key objects (e.g., vehicles, weapons, bags) in the image. Provide object names and approximate bounding box coordinates (x_min, y_min, x_max, y_max) in the image."
                inputs_objects = processor(text=[prompt_objects], images=[image], return_tensors="pt", padding=True)
                objects_outputs = model.generate(**inputs_objects, max_new_tokens=256)
                result["objects"] = [{"object": obj.strip(), "bbox": "(unknown)"} for obj in processor.decode(objects_outputs[0], skip_special_tokens=True).replace(prompt_objects, "").split(",") if obj.strip()] or [{"object": "None", "bbox": "N/A"}]
            
            if activity:
                if cancel_flag:
                    return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
                prompt_activity = "<image> Describe any activities or actions occurring in the image, such as walking, running, or driving."
                inputs_activity = processor(text=[prompt_activity], images=[image], return_tensors="pt", padding=True)
                activity_outputs = model.generate(**inputs_activity, max_new_tokens=256)
                result["activity"] = processor.decode(activity_outputs[0], skip_special_tokens=True).replace(prompt_activity, "").strip() or "No activity detected"
            
            # Investigation-Specific Features
            if clothing:
                if cancel_flag:
                    return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
                prompt_clothing = "<image> Identify clothing items and their colors worn by people in the image."
                inputs_clothing = processor(text=[prompt_clothing], images=[image], return_tensors="pt", padding=True)
                clothing_outputs = model.generate(**inputs_clothing, max_new_tokens=256)
                result["clothing"] = processor.decode(clothing_outputs[0], skip_special_tokens=True).replace(prompt_clothing, "").strip() or "No clothing detected"
            
            if scene:
                if cancel_flag:
                    return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
                prompt_scene = "<image> Classify the scene type (e.g., indoor, outdoor, urban, rural) and estimate the time of day (e.g., day, night, dusk)."
                inputs_scene = processor(text=[prompt_scene], images=[image], return_tensors="pt", padding=True)
                scene_outputs = model.generate(**inputs_scene, max_new_tokens=256)
                result["scene_context"] = processor.decode(scene_outputs[0], skip_special_tokens=True).replace(prompt_scene, "").strip() or "No context determined"
            
            if threat_score:
                if cancel_flag:
                    return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
                result["threat_score"] = calculate_threat_score(result.get("harmful_objects", []), weights)
            
            if similarity:
                if cancel_flag:
                    return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
                inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
                with torch.no_grad():
                    emb = model.vision_model(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
                image_embeddings.append(emb)
            
            results.append(result)
        except Exception as e:
            error_result = {
                "image_id": f"Image_{idx}",
                "timestamp": timestamp,
                "error": str(e)
            }
            error_result.update({
                "description": "Error processing image." if description else None,
                "signs": "Error" if signs else None,
                "harmful_objects": [{"object": "Error", "confidence": 0}] if harmful else None,
                "faces": "Error" if faces else None,
                "objects": [{"object": "Error", "bbox": "N/A"}] if objects else None,
                "scene_context": "Error" if scene else None,
                "activity": "Error" if activity else None,
                "clothing": "Error" if clothing else None,
                "metadata": {"error": str(e)} if metadata else None,
                "threat_score": 0.0 if threat_score else None
            })
            error_result = {k: v for k, v in error_result.items() if v is not None}
            results.append(error_result)
    
    # Compute similarity scores
    if similarity and len(image_embeddings) > 1:
        base_embedding = image_embeddings[0]
        for i in range(1, len(image_embeddings)):
            if cancel_flag:
                return {"results": results, "analysis_timestamp": timestamp, "status": "Cancelled"}
            sim = np.dot(base_embedding, image_embeddings[i].T) / (
                np.linalg.norm(base_embedding) * np.linalg.norm(image_embeddings[i])
            )
            results[i]["similarity_to_image_1"] = float(sim[0][0])
    
    return {"results": results, "analysis_timestamp": timestamp, "status": "Completed"}

# Gradio interface
def gradio_predict(
    image_1, image_2, image_3,
    description: bool,
    signs: bool,
    harmful: bool,
    similarity: bool,
    faces: bool,
    objects: bool,
    scene: bool,
    metadata: bool,
    activity: bool,
    clothing: bool,
    threat_score: bool,
    combined: bool,
    json_export: bool,
    detailed_report: bool,
    custom_harmful: str,
    custom_weights: str,
    description_level: str,
    keyword_search: str,
    filter_attributes: str,
    annotation: str,
    flag_images: bool,
    comments: str,
    case_folder: str,
    case_action: str
):
    global cancel_flag
    cancel_flag = False  # Reset cancellation flag
    images = [image_1, image_2, image_3]
    images = [img for img in images if img is not None]
    if not images:
        return "Error: At least one image must be uploaded."
    if len(images) > 3:
        return "Error: Maximum 3 images allowed."
    if not any([description, signs, harmful, similarity, faces, objects, scene, metadata, activity, clothing, threat_score, combined]):
        return "Error: At least one output option must be selected."
    
    # If combined is selected, enable all outputs
    if combined:
        description = signs = harmful = similarity = faces = objects = scene = metadata = activity = clothing = threat_score = True
    
    # Process custom harmful objects and weights
    harmful_objects = default_harmful_objects.copy()
    try:
        custom_objects = validate_custom_harmful(custom_harmful)
        harmful_objects.extend(custom_objects)
    except ValueError as e:
        return f"Error: {str(e)}"
    
    weights = {obj: 1.0 for obj in harmful_objects}
    if custom_weights.strip():
        try:
            for pair in custom_weights.split(","):
                obj, weight = pair.split(":")
                weights[obj.strip().lower()] = float(weight.strip())
        except:
            return "Error: Invalid custom weights format. Use 'object:weight,object:weight'."
    
    results = []
    image_embeddings = []
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    errors = []
    
    # Precompute key findings
    findings = {
        "harmful_found": False,
        "high_similarity": False,
        "faces_detected": False,
        "threats_detected": False
    }
    
    for idx, image in enumerate(images, 1):
        if cancel_flag:
            output = f"**Analysis Cancelled at**: {timestamp}\n"
            if results:
                output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
            return output
        try:
            # Convert Gradio image input to PIL and resize
            image = Image.fromarray(image).convert("RGB")
            image = resize_image(image, max_size=1024)
            image_data = io.BytesIO()
            image.save(image_data, format="JPEG")
            image_data.seek(0)
            
            result = {"image_id": f"Image_{idx}", "timestamp": timestamp}
            
            # Core Features
            if metadata:
                if cancel_flag:
                    output = f"**Analysis Cancelled at**: {timestamp}\n"
                    if results:
                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                    return output
                result["metadata"] = extract_metadata(image_data.getvalue())
            
            if description:
                if cancel_flag:
                    output = f"**Analysis Cancelled at**: {timestamp}\n"
                    if results:
                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                    return output
                prompt_desc = "<image> Provide a " + ("brief description of the image." if description_level == "basic" else "detailed description of the image, including objects, colors, people, and environmental context.")
                inputs = processor(text=[prompt_desc], images=[image], return_tensors="pt", padding=True)
                outputs = model.generate(**inputs, max_new_tokens=256 if description_level == "basic" else 512)
                result["description"] = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip() or "No description generated."
            
            if signs:
                if cancel_flag:
                    output = f"**Analysis Cancelled at**: {timestamp}\n"
                    if results:
                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                    return output
                prompt_ocr = "<image> Extract all visible text in the image, such as road signs, license plates, or billboards, with exact wording."
                inputs_ocr = processor(text=[prompt_ocr], images=[image], return_tensors="pt", padding=True)
                ocr_outputs = model.generate(**inputs_ocr, max_new_tokens=256)
                result["signs"] = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip() or "None detected"
            
            if harmful:
                if cancel_flag:
                    output = f"**Analysis Cancelled at**: {timestamp}\n"
                    if results:
                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                    return output
                prompt_detect = f"<image> Identify any harmful objects ({', '.join(harmful_objects)}) in the image. List them explicitly and estimate confidence (0-100%) for each detection based on clarity."
                inputs_detect = processor(text=[prompt_detect], images=[image], return_tensors="pt", padding=True)
                detect_outputs = model.generate(**inputs_detect, max_new_tokens=256)
                detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
                harmful_detected = []
                for obj in harmful_objects:
                    if obj in detected_objects:
                        confidence = 90 if obj in detected_objects.split() else 60
                        harmful_detected.append({"object": obj, "confidence": confidence})
                result["harmful_objects"] = harmful_detected if harmful_detected else [{"object": "None", "confidence": 0}]
                if harmful_detected and result["harmful_objects"][0]["object"] != "None":
                    findings["harmful_found"] = True
            
            if faces:
                if cancel_flag:
                    output = f"**Analysis Cancelled at**: {timestamp}\n"
                    if results:
                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                    return output
                prompt_faces = "<image> Detect faces and estimate attributes such as age range (e.g., child, adult, senior), gender (male, female, unknown), and emotional cues (e.g., neutral, angry, scared)."
                inputs_faces = processor(text=[prompt_faces], images=[image], return_tensors="pt", padding=True)
                faces_outputs = model.generate(**inputs_faces, max_new_tokens=256)
                result["faces"] = processor.decode(faces_outputs[0], skip_special_tokens=True).replace(prompt_faces, "").strip() or "No faces detected"
                if result["faces"] != "No faces detected":
                    findings["faces_detected"] = True
            
            if objects:
                if cancel_flag:
                    output = f"**Analysis Cancelled at**: {timestamp}\n"
                    if results:
                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                    return output
                prompt_objects = "<image> Identify and localize key objects (e.g., vehicles, weapons, bags) in the image. Provide object names and approximate bounding box coordinates (x_min, y_min, x_max, y_max) in the image."
                inputs_objects = processor(text=[prompt_objects], images=[image], return_tensors="pt", padding=True)
                objects_outputs = model.generate(**inputs_objects, max_new_tokens=256)
                result["objects"] = [{"object": obj.strip(), "bbox": "(unknown)"} for obj in processor.decode(objects_outputs[0], skip_special_tokens=True).replace(prompt_objects, "").split(",") if obj.strip()] or [{"object": "None", "bbox": "N/A"}]
            
            if activity:
                if cancel_flag:
                    output = f"**Analysis Cancelled at**: {timestamp}\n"
                    if results:
                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                    return output
                prompt_activity = "<image> Describe any activities or actions occurring in the image, such as walking, running, or driving."
                inputs_activity = processor(text=[prompt_activity], images=[image], return_tensors="pt", padding=True)
                activity_outputs = model.generate(**inputs_activity, max_new_tokens=256)
                result["activity"] = processor.decode(activity_outputs[0], skip_special_tokens=True).replace(prompt_activity, "").strip() or "No activity detected"
            
            # Investigation-Specific Features
            if clothing:
                if cancel_flag:
                    output = f"**Analysis Cancelled at**: {timestamp}\n"
                    if results:
                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                    return output
                prompt_clothing = "<image> Identify clothing items and their colors worn by people in the image."
                inputs_clothing = processor(text=[prompt_clothing], images=[image], return_tensors="pt", padding=True)
                clothing_outputs = model.generate(**inputs_clothing, max_new_tokens=256)
                result["clothing"] = processor.decode(clothing_outputs[0], skip_special_tokens=True).replace(prompt_clothing, "").strip() or "No clothing detected"
            
            if scene:
                if cancel_flag:
                    output = f"**Analysis Cancelled at**: {timestamp}\n"
                    if results:
                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                    return output
                prompt_scene = "<image> Classify the scene type (e.g., indoor, outdoor, urban, rural) and estimate the time of day (e.g., day, night, dusk)."
                inputs_scene = processor(text=[prompt_scene], images=[image], return_tensors="pt", padding=True)
                scene_outputs = model.generate(**inputs_scene, max_new_tokens=256)
                result["scene_context"] = processor.decode(scene_outputs[0], skip_special_tokens=True).replace(prompt_scene, "").strip() or "No context determined"
            
            if threat_score:
                if cancel_flag:
                    output = f"**Analysis Cancelled at**: {timestamp}\n"
                    if results:
                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                    return output
                result["threat_score"] = calculate_threat_score(result.get("harmful_objects", []), weights)
                if result["threat_score"] > 50:
                    findings["threats_detected"] = True
            
            # User Options & Controls
            if annotation and annotation.strip():
                result["annotation"] = annotation.strip()
            if flag_images:
                result["flag"] = True
            if comments and comments.strip():
                result["comments"] = comments.strip()
            
            if similarity:
                if cancel_flag:
                    output = f"**Analysis Cancelled at**: {timestamp}\n"
                    if results:
                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                    return output
                inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
                with torch.no_grad():
                    emb = model.vision_model(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
                image_embeddings.append(emb)
            
            results.append(result)
        except Exception as e:
            errors.append({"image_id": f"Image_{idx}", "error": str(e)})
            result = {
                "image_id": f"Image_{idx}",
                "timestamp": timestamp,
                "error": str(e)
            }
            result.update({
                "description": "Error processing image." if description else None,
                "signs": "Error" if signs else None,
                "harmful_objects": [{"object": "Error", "confidence": 0}] if harmful else None,
                "faces": "Error" if faces else None,
                "objects": [{"object": "Error", "bbox": "N/A"}] if objects else None,
                "scene_context": "Error" if scene else None,
                "activity": "Error" if activity else None,
                "clothing": "Error" if clothing else None,
                "metadata": {"error": str(e)} if metadata else None,
                "threat_score": 0.0 if threat_score else None
            })
            result = {k: v for k, v in result.items() if v is not None}
            if annotation and annotation.strip():
                result["annotation"] = annotation.strip()
            if flag_images:
                result["flag"] = True
            if comments and comments.strip():
                result["comments"] = comments.strip()
            results.append(result)
    
    # Compute similarity scores
    if similarity and len(image_embeddings) > 1:
        base_embedding = image_embeddings[0]
        for i in range(1, len(image_embeddings)):
            if cancel_flag:
                output = f"**Analysis Cancelled at**: {timestamp}\n"
                if results:
                    output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                return output
            sim = np.dot(base_embedding, image_embeddings[i].T) / (
                np.linalg.norm(base_embedding) * np.linalg.norm(image_embeddings[i])
            )
            sim_value = float(sim[0][0])
            results[i]["similarity_to_image_1"] = sim_value
            if sim_value > 0.8:
                findings["high_similarity"] = True
    
    # Filter by attributes
    if filter_attributes and filter_attributes.strip():
        try:
            attributes = [attr.strip().lower() for attr in filter_attributes.split(",") if attr.strip()]
            results = [
                result for result in results
                if any(
                    any(attr in str(value).lower() for value in result.values() if value is not None)
                    for attr in attributes
                )
            ]
        except:
            errors.append({"error": "Invalid filter attributes format."})
    
    # Handle case folder
    case_output = ""
    if case_folder and case_folder.strip() and not cancel_flag:
        case_output = manage_case_folder(case_folder, results, case_action)
    
    # Format output
    output = format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
    
    # Reset cancellation flag
    cancel_flag = False
    return output

# Helper function to format results
def format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors):
    output = [f"**Analysis Timestamp**: {timestamp}\n\n"]
    headers = ["Image ID"]
    if metadata:
        headers.append("Metadata")
    if description:
        headers.append("Description")
    if signs:
        headers.append("Signs")
    if harmful:
        headers.append("Harmful Objects")
    if faces:
        headers.append("Faces")
    if objects:
        headers.append("Objects")
    if clothing:
        headers.append("Clothing")
    if scene:
        headers.append("Scene")
    if activity:
        headers.append("Activity")
    if threat_score:
        headers.append("Threat Score")
    if similarity:
        headers.append("Similarity")
    if annotation and annotation.strip():
        headers.append("Annotation")
    if flag_images:
        headers.append("Flagged")
    if comments and comments.strip():
        headers.append("Comments")
    
    output.append("| " + " | ".join(headers) + " |\n")
    output.append("| " + " | ".join(["---"] * len(headers)) + " |\n")
    
    filtered_results = keyword_search(results, keyword_search)
    for result in filtered_results or []:
        row = [result.get('image_id', 'N/A')]
        if metadata:
            meta = result.get("metadata", {})
            meta_str = f"Time: {meta.get('timestamp', 'N/A')}, GPS: {meta.get('gps', {}).get('latitude', 'N/A')}, {meta.get('gps', {}).get('longitude', 'N/A')}"
            row.append(meta_str if "error" not in meta else meta.get("error", "N/A"))
        if description:
            row.append(result.get("description", "N/A"))
        if signs:
            row.append(result.get("signs", "N/A"))
        if harmful:
            harmful_str = ", ".join([f"{obj.get('object', 'N/A')} ({obj.get('confidence', 0)}%)" for obj in result.get("harmful_objects", [])]) or "None"
            row.append(harmful_str)
        if faces:
            row.append(result.get("faces", "N/A"))
        if objects:
            objects_str = ", ".join([f"{obj.get('object', 'N/A')} at {obj.get('bbox', 'N/A')}" for obj in result.get("objects", [])]) or "None"
            row.append(objects_str)
        if clothing:
            row.append(result.get("clothing", "N/A"))
        if scene:
            row.append(result.get("scene_context", "N/A"))
        if activity:
            row.append(result.get("activity", "N/A"))
        if threat_score:
            row.append(f"{result.get('threat_score', 0):.1f}" if "threat_score" in result else "N/A")
        if similarity:
            similarity_val = f"{result.get('similarity_to_image_1', 0):.2f}" if 'similarity_to_image_1' in result else "N/A"
            row.append(similarity_val)
        if annotation and annotation.strip():
            row.append(result.get("annotation", "N/A"))
        if flag_images:
            row.append("Yes" if result.get("flag", False) else "No")
        if comments and comments.strip():
            row.append(result.get("comments", "N/A"))
        
        output.append("| " + " | ".join(row) + " |\n")
    
    # Append errors
    if errors:
        output.append("\n**Errors**:\n")
        for error in errors:
            output.append(f"- {error.get('image_id', 'N/A')}: {error.get('error', 'Unknown error')}\n")
    
    # Handle JSON export
    if json_export:
        json_output = {"results": filtered_results, "analysis_timestamp": timestamp}
        output.append("\n**JSON Export**:\n```json\n" + json.dumps(json_output, indent=2, default=str) + "\n```")
    
    # Handle detailed report
    if detailed_report:
        output.append("\n**Detailed Report**:\n" + format_detailed_report(results, timestamp, findings, keyword_search))
    
    # Append case folder output
    if case_folder and case_folder.strip() and case_output:
        output.append("\n**Case Folder**:\n" + case_output)
    
    return "".join(output)

# Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# VisionSage: Image Analysis for Crime Investigation")
    gr.Markdown("Upload up to 3 images (up to 10MB each, any resolution) ")
    
    with gr.Row():
        image_1 = gr.Image(label="Upload Image 1 (up to 10MB)")
        image_2 = gr.Image(label="Upload Image 2 (up to 10MB)")
        image_3 = gr.Image(label="Upload Image 3 (up to 10MB)")
    
    with gr.Row():
        with gr.Column():
            description = gr.Checkbox(label="Description", value=True)
            signs = gr.Checkbox(label="Signs", value=True)
            harmful = gr.Checkbox(label="Harmful Objects", value=True)
            similarity = gr.Checkbox(label="Similarity", value=True)
            faces = gr.Checkbox(label="Facial Attributes", value=False)
            objects = gr.Checkbox(label="Localized Objects", value=False)
            scene = gr.Checkbox(label="Scene Context", value=False)
            metadata = gr.Checkbox(label="Metadata", value=False)
            activity = gr.Checkbox(label="Activity Recognition", value=False)
            clothing = gr.Checkbox(label="Clothing/Colors", value=False)
            threat_score = gr.Checkbox(label="Threat Score", value=False)
            combined = gr.Checkbox(label="Combined (All Outputs)", value=False)
        
        with gr.Column():
            json_export = gr.Checkbox(label="JSON Export", value=False)
            detailed_report = gr.Checkbox(label="Detailed Report", value=False)
            custom_harmful = gr.Textbox(label="Custom Harmful Objects (comma-separated, e.g., handgun, crowbar)", placeholder="Enter objects to detect")
            custom_weights = gr.Textbox(label="Custom Weights (e.g., knife:2.0,gun:3.0)", placeholder="Enter object:weight pairs")
            description_level = gr.Radio(label="Description Level", choices=["basic", "detailed"], value="detailed")
            keyword_search = gr.Textbox(label="Keyword Search", placeholder="Enter keywords to filter results")
            filter_attributes = gr.Textbox(label="Filter by Attributes (comma-separated)", placeholder="e.g., red, adult, urban")
            annotation = gr.Textbox(label="Manual Annotation", placeholder="Add labels or notes to images")
            flag_images = gr.Checkbox(label="Flag Important Images", value=False)
            comments = gr.Textbox(label="Investigator Comments", placeholder="Add comments or notes")
            case_folder = gr.Textbox(label="Case Folder Name", placeholder="Enter case folder name")
            case_action = gr.Radio(label="Case Folder Action", choices=["add", "view", "clear"], value="add")
    
    with gr.Row():
        submit_button = gr.Button("Submit")
        cancel_button = gr.Button("Cancel Analysis")
    
    output = gr.Textbox(label="Investigation Results", placeholder="Results will appear here...")
    
    submit_button.click(
        fn=gradio_predict,
        inputs=[
            image_1, image_2, image_3,
            description, signs, harmful, similarity, faces, objects, scene, metadata, activity, clothing, threat_score, combined,
            json_export, detailed_report, custom_harmful, custom_weights, description_level, keyword_search, filter_attributes,
            annotation, flag_images, comments, case_folder, case_action
        ],
        outputs=output
    )
    
    cancel_button.click(
        fn=cancel_analysis,
        inputs=[],
        outputs=output
    )

if __name__ == "__main__":
    # Launch Gradio interface for Hugging Face Spaces
    iface.launch(server_name="0.0.0.0", server_port=7860)