# === FASTAPI BACKEND (main.py) ===

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline
from PIL import Image
import io
import torch
import numpy as np
import cv2
import base64

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

device = 0 if torch.cuda.is_available() else -1

MODELS_CONFIG = {
    "SwinV2 Based": {"path": "haywoodsloan/ai-image-detector-deploy", "weight": 0.15},
    "ViT Based": {"path": "Heem2/AI-vs-Real-Image-Detection", "weight": 0.15},
    "SDXL Dataset": {"path": "Organika/sdxl-detector", "weight": 0.15},
    "SDXL + FLUX": {"path": "cmckinle/sdxl-flux-detector_v1.1", "weight": 0.15},
    "DeepFake v2": {"path": "prithivMLmods/Deep-Fake-Detector-v2-Model", "weight": 0.15},
    "Midjourney/SDXL": {"path": "ideepankarsharma2003/AI_ImageClassification_MidjourneyV6_SDXL", "weight": 0.10},
    "ViT v4": {"path": "date3k2/vit-real-fake-classification-v4", "weight": 0.15},
}

models = {}
for name, config in MODELS_CONFIG.items():
    try:
        models[name] = pipeline("image-classification", model=config["path"], device=device)
    except Exception as e:
        print(f"Failed to load model {name}: {e}")

def pil_to_base64(image):
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    return "data:image/jpeg;base64," + base64.b64encode(buffered.getvalue()).decode("utf-8")

def gen_ela(img_array, quality=90):
    if img_array.shape[2] == 4:
        img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
    _, buffer = cv2.imencode('.jpg', img_array, encode_param)
    compressed_img = cv2.imdecode(buffer, cv2.IMREAD_COLOR)
    ela_img = cv2.absdiff(img_array, compressed_img)
    ela_img = cv2.convertScaleAbs(ela_img, alpha=10)
    return Image.fromarray(cv2.cvtColor(ela_img, cv2.COLOR_BGR2RGB))

def gradient_processing(image_array):
    gray_img = cv2.cvtColor(image_array, cv2.COLOR_BGR2GRAY)
    dx = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=3)
    dy = cv2.Sobel(gray_img, cv2.CV_64F, 0, 1, ksize=3)
    gradient_magnitude = cv2.magnitude(dx, dy)
    gradient_img = cv2.normalize(gradient_magnitude, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    return Image.fromarray(gradient_img)

@app.post("/detect")
async def detect(image: UploadFile = File(...)):
    try:
        import time
        start_time = time.time()

        image_bytes = await image.read()
        input_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

        img_np = np.array(input_image)
        img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)

        individual_results = []
        weighted_ai_score = 0
        total_weight = 0

        aiModels = []
        colors = ["bg-red-500", "bg-orange-500", "bg-yellow-500", "bg-green-500", "bg-blue-500", "bg-purple-500", "bg-pink-500"]

        for i, (name, model_pipeline) in enumerate(models.items()):
            model_weight = MODELS_CONFIG[name]["weight"]
            predictions = model_pipeline(input_image)
            confidence = {p['label'].lower(): p['score'] for p in predictions}

            artificial_score = (
                confidence.get('artificial', 0) or confidence.get('ai image', 0) or
                confidence.get('ai', 0) or confidence.get('deepfake', 0) or
                confidence.get('ai_gen', 0) or confidence.get('fake', 0)
            )
            real_score = (
                confidence.get('real', 0) or confidence.get('real image', 0) or
                confidence.get('human', 0) or confidence.get('realism', 0)
            )

            if artificial_score > 0 and real_score == 0:
                real_score = 1.0 - artificial_score
            elif real_score > 0 and artificial_score == 0:
                artificial_score = 1.0 - real_score

            weighted_ai_score += artificial_score * model_weight
            total_weight += model_weight

            aiModels.append({
                "name": name,
                "percentage": round(artificial_score * 100, 2),
                "color": colors[i % len(colors)]
            })

        final_score = (weighted_ai_score / total_weight) * 100 if total_weight > 0 else 0
        verdict = final_score > 50
        processing_time = int((time.time() - start_time) * 1000)

        # Forensics
        ela_img = gen_ela(img_bgr)
        gradient_img = gradient_processing(img_bgr)

        return JSONResponse({
            "filename": image.filename,
            "isDeepfake": verdict,
            "confidence": round(final_score, 2),
            "aiModels": aiModels,
            "processingTime": processing_time,
            "forensics": {
                "original": pil_to_base64(input_image),
                "ela": pil_to_base64(ela_img),
                "gradient": pil_to_base64(gradient_img)
            },
            "verdictMessage": f"Consensus: {'Likely AI-Generated' if verdict else 'Likely Human-Made (Real)'}"
        })
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))