import os
import gradio as gr
import warnings
from huggingface_hub import Repository, login

warnings.filterwarnings("ignore")

login(token=os.environ['HF_TOKEN'])

repo = Repository(
    local_dir="dataset",
    repo_type="dataset",
    clone_from=os.environ['DATASET'],
    token=True
)
repo.git_pull()

# Import the AI detector
from dataset.model import AIContentDetector

# Initialize the detector
detector = AIContentDetector()

def analyze_image(image):
    """Wrapper function for image analysis."""
    if image is None:
        return "Please upload an image", 0
    
    result, confidence = detector.detect_ai_image(image)
    return result, confidence

def analyze_video(video):
    """Wrapper function for video analysis."""
    if video is None:
        return "Please upload a video", 0
    
    result, confidence = detector.detect_ai_video(video)
    return result, confidence

def analyze_audio(audio):
    """Wrapper function for audio analysis."""
    if audio is None:
        return "Please upload an audio file", 0
    
    result, confidence = detector.detect_ai_audio(audio)
    return result, confidence

# Create the Gradio interface
demo = gr.Interface(
    fn=analyze_image,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Textbox(label="Detection Result"),
        gr.Number(label="Confidence Score (%)")
    ],
    title="🕵️ GuanFu - AI Content Detector",
    description="Upload an image to detect if it was generated by AI.",
    theme=gr.themes.Soft()
)

if __name__ == "__main__":
    demo.launch()