import os
import json
import base64
import requests
import gradio as gr
from dotenv import load_dotenv

# Load environment variables (useful for local testing)
load_dotenv()

# The API Key must be set in the HuggingFace Space "Secrets"
API_KEY = os.getenv("HIVE_API_KEY")
HIVE_API_URL = "https://api.thehive.ai/api/v3/hive/ai-generated-and-deepfake-content-detection"

def detect_ai_content(image_filepath):
    if not API_KEY:
        raise gr.Error("HIVE_API_KEY environment variable is not set. Please add it to your HuggingFace Space Secrets.")
        
    if image_filepath is None:
        raise gr.Error("No image provided.")

    try:
        # Read the image file and encode it to Base64
        with open(image_filepath, "rb") as image_file:
            file_bytes = image_file.read()
        
        encoded_string = base64.b64encode(file_bytes).decode("utf-8")
        
        # Format the Base64 string for the Hive API
        media_base64 = f"data:image/jpeg;base64,{encoded_string}"

        payload = {
            "media_metadata": True,
            "input": [{"media_base64": media_base64}]
        }

        headers = {
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json",
            "Accept": "application/json"
        }

        response = requests.post(
            HIVE_API_URL,
            headers=headers,
            data=json.dumps(payload),
            timeout=60
        )
        
        # Raise an exception if the response indicates an HTTP error
        response.raise_for_status()
        
        data = response.json()
        
        real_val = 0.0
        fake_val = 0.0
        
        # Parse the JSON to extract the required classes
        if "output" in data and len(data["output"]) > 0:
            classes = data["output"][0].get("classes", [])
            for c in classes:
                class_name = c.get("class", "")
                if class_name == "not_ai_generated" or class_name == "none":
                    real_val = c.get("value", 0.0)
                elif class_name == "ai_generated":
                    fake_val = c.get("value", 0.0)
            
            # If the API returned specific model percentages (like midjourney, dall-e) 
            # but no overarching 'ai_generated' tag, we sum all fake probabilities
            if fake_val == 0.0 and real_val > 0.0:
                fake_val = sum(c.get("value", 0.0) for c in classes if c.get("class") not in ["not_ai_generated", "none", "inconclusive", "not_ai_generated_audio", "ai_generated_audio"])
                    
        return {"Real": real_val, "Fake": fake_val}
        
    except requests.exceptions.RequestException as e:
        error_msg = f"API Error: {e}"
        if hasattr(e, 'response') and e.response is not None:
            error_msg += f" | Response: {e.response.text}"
        raise gr.Error(error_msg)
    except Exception as e:
        raise gr.Error(f"Internal Error: {e}")

# Define the Gradio Interface
demo = gr.Interface(
    fn=detect_ai_content,
    inputs=gr.Image(type="filepath", label="Upload an Image"),
    outputs=gr.Label(label="Detection Results"),
    title="Hive AI Content Detection",
    description="Upload an image to detect if it is AI-generated or contains deepfakes using the Hive AI API."
)

# Launch the app
if __name__ == "__main__":
    demo.launch()