# server.py - DOUBLE FALLBACK ARCHITECTURE (SAPLING -> 7D | SIGHTENGINE -> CONVNEXT)
import os
import io
import gc
import cv2
import math
import uuid
import shutil
import joblib
import zipfile
import numpy as np
import torch
import torch.nn.functional as F
import timm
from collections import Counter
from typing import Optional
import requests  

# API & Image Handling
from fastapi import FastAPI, HTTPException, UploadFile, File, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from PIL import Image
from torchvision import transforms
from skimage.measure import shannon_entropy
from docx import Document
import nltk

# Transformers & Hub
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from huggingface_hub import hf_hub_download, list_repo_files

# ==========================================================
# 0. NLTK SETUP
# ==========================================================
try:
    nltk.data.find('tokenizers/punkt')
    nltk.data.find('tokenizers/punkt_tab')
except LookupError:
    nltk.download('punkt')
    nltk.download('punkt_tab')

# ==========================================================
# 1. CONFIGURATION & EXTERNAL APIs
# ==========================================================

# --- TEXT API: Sapling ---
SAPLING_URL = "https://api.sapling.ai/api/v1/aidetect"
SAPLING_API_KEYS = [
    "LFKP2ESNCF82CCX07XBOIY26ZMFT4QEO", 
    "4ZB9X067IWMVXRTJLKNVJKVDWQZEXHZ8",     
    "AECPJC27LX68E5L6O5X2V75S95IWUWKK",
    "S5D99OJYA6TXRX49JJWKQUJP8P3YWYTH",
    "QDB1FMQQZXOUJYRISG1G0K7QJGMPFTT4",
    "YUUPMELNIJ424Z9H8F3NSG3529Q1VB86",
    "85DNYDUMEASIHNQLSTHRCBJCN5B7S8C9",
    "AT74JLK6VSJRR5O033OR73T39ABUQG6E",
    "K1RG0FXGPQYXXMSM2K9BFISYJFZ5ZMWQ",
    "ATXVUCTJCRQKWUG2P5S15A9EEA0BPAW3",
    "57ND7C0EOHMC6YW9B89T0BYLYXVTO5VG",
    "XAZKMDU879F6KX2RWSTHRZ1X0XLQD5BK",
    "KA91IC2RKH68Z78CLJ4YDGDYU3BK9FIF",
    "SPS0WY0377ZRIM3IKATC7MX5G8XDPFU8",
    "NXTFWMCK47C7N0SMVK2HTGRBN00608E4",
    "C8UX45IZMPADBFYR8FWKN9GB92TODQO6",
    "MOMKK9QJ9W3SEDILWPMKSEH36EBK7XI6",
    "YYYUILP6BWHG6XD83QQ0Y3IJT431VOD6",
    "FVZ9FBFYFP0E93DHFDRH75V82F1TIFPW",
    "YJ99AR2XVVR581LX4W93KCW59ZNQOW9O",
    "47YR5OFIGWT3N9LS2VUL6Z9MY32XG3B1",
    "A9YHY7ARNOU46AJOYBYM3HZGGKRHTKLX",
    "86S3VGGQMCEF54R0IBEPT3N8RVLU6CFU",
    "0VGCCYKM9H5PWTPHT7YOK1GGUDV77MO3",
    "XU45624B3O6PEEMAT7BY7FNQB3U5W6MR",
    "TKGDY124QDM20ERR67MP171GJ8KEVU07",
    "V0G6NYIE13S8C8G0KOBI8I7GJLP8TPUX",
    "76MB73YTW8N4XV3N66ZC3Y6XJL9BEAJ6",
    "E2AIR6LCQO38PHTHXRQSOBP0P8Z28X58",
    "U4ZIOTZ1K72TB6V5VXU9BLZE3NQYPTVS",
    "PRZQS3EFQZ9253Q62VUTTUKRVJH5UJ5W",
    "N5VU02OZHXHEAGT7AUKNI41UD9GJB8LY",
]
current_sapling_idx = 0

# --- IMAGE API: Sightengine ---
SIGHTENGINE_URL = "https://api.sightengine.com/1.0/check.json"
SIGHTENGINE_API_KEYS = [
    {"user": "1743982661", "secret": "ZEZwGjcsyvq6ULyKgj5456SzNscWtKB8"},
    {"user": "1638959208", "secret": "sei2oE3bbvzXsm5Nfs8tujdpbXCqETjb"},
    {"user": "1882880449", "secret": "LyHJkg6qWAEztt8CVrPsyNRxR2y97H7s"},
    {"user": "1776948743", "secret": "CXVznzPQU9P55AduTna4LggSrpHDYCrL"},
    {"user": "1168060613", "secret": "8wRJKanHWkPLjnhDKb6GspdzMRcoiB7X"},
    {"user": "20376699", "secret": "ruavg86UNL3wn8chGKidNX7spYDAJgrL"},
    {"user": "1525731149", "secret": "JfjjCJJzdyGWbaxWioHWDGGqPSWkskKc"},
    {"user": "482807524", "secret": "sFjXFDyBTiuyh33C57jWicfLVsPGh35W"},
    {"user": "1806137219", "secret": "HcXtBp7rXBCBZExQzXRnPWo7HrAx7Mon"},
    {"user": "1705951192", "secret": "9W7xT9AUiKJjeCiG6Q5CjT7zYHdQefJB"},
    {"user": "1632737246", "secret": "PBEeTJVPXAJMhYUgnebzm4p2RrGh4BQz"},
    {"user": "857335768", "secret": "bAqSCXTBRHUEy8E3vdH42JXf2rABpVKo"},
    {"user": "1147336463", "secret": "2AKvLQN4VW33QDGvq6tY598SRwRYhGHu"},
    {"user": "820176459", "secret": "F3SGYgiztjjCSMJo2DV4yHVWEpeTLzzx"},
    {"user": "493725389", "secret": "HmsjtxMp9CSzfASkWTtK3zFYHPXaF6zK"},
    {"user": "91559601", "secret": "F2KMUdDPq322fT9aDkPmbZjgjqxSHqef"},
    {"user": "1478152006", "secret": "bVJ6xtikepzxrVcKKPRdMKL4h7aWbeek"},
    {"user": "361978076", "secret": "ShjdhXtXuk9jMjGBMek4hDZuyZXheNNF"},
    {"user": "355220504", "secret": "kBcuc7RSYoAqxqmQkzwFUbH5wvBD8Gwr"},
    {"user": "1524398481", "secret": "5mZzcBiSdmf5KMSrUDXFojU6zJdNUeKr"},
    {"user": "1751980311", "secret": "JuCLmEDBQvYw6ojDHiHgpm2YuRPWxCH3"},
    {"user": "411151581", "secret": "Dm6GNEQtWMMDCPKY6wAb8X73x8BxeSg6"},
    {"user": "576637210", "secret": "JTYDsX9iDZdifKoZ9uzhbMUSENVazFqs"},
    {"user": "106682960", "secret": "8J6EuKLN5Pue6smijuBdqqAQX3kaYbGd"},
    {"user": "1826450561", "secret": "jqnGBwYSHnyj8jsvwc6gA7yFbU3UpX5T"},
]
current_sightengine_idx = 0

# --- Models ---
TEXT_MODEL_1_ID = "Yuvrajg2107/deberta-v3-hybrid-detector_v2_universal"
TEXT_MODEL_2_ID = "Yuvrajg2107/roberta-base-cpp-final"
TEXT_MODEL_3_ID = "Yuvrajg2107/electra-large-discriminator-cpp-final"
ENSEMBLE_PATH = "model_ensemble_pro.pkl"

MARATHI_MAHABERT_ID = "Yashodhar29/mahabert-marathi-detection"
MARATHI_INDICBERT_ID = "Yashodhar29/indicbert-marathi-detection"
MARATHI_ENSEMBLE_PATH = "marathi_ensemble_pro.pkl"
CODE_MODEL_ID = "Yashodhar29/Qwen2.5-Coder-0.5B-Instruct-cpp"

# The Fallback Image Model
IMAGE_REPO_ID = "Yashodhar29/ConvNext-large-cpp"

# --- Device Setup ---
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🚀 Server starting on device: {device.upper()}")

# ==========================================================
# 2. MODEL LOADING INFRASTRUCTURE
# ==========================================================

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

# Global Model Storage
models = {
    "text": [],
    "code": None,
    "marathi": None,
    "image": None # Local ConvNeXt Fallback
}

def load_text_model(model_id):
    print(f"   ⏳ Loading {model_id}...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        model = AutoModelForSequenceClassification.from_pretrained(
            model_id, 
            trust_remote_code=True,
            low_cpu_mem_usage=True
        ).to(device)
        model.eval()
        return {"model": model, "tokenizer": tokenizer, "name": model_id}
    except Exception as e:
        print(f"   ❌ Failed to load {model_id}: {e}")
        return None

def load_image_model_from_hub(repo_id):
    print(f"   ⏳ Loading Fallback Image Repo: {repo_id}...")
    try:
        files = list_repo_files(repo_id)
        pth_files = [f for f in files if f.endswith('.pth')]
        if not pth_files: return None
        
        weights_path = hf_hub_download(repo_id=repo_id, filename=pth_files[0])
        model = timm.create_model("convnext_large.fb_in22k_ft_in1k", pretrained=False, num_classes=2)
        state_dict = torch.load(weights_path, map_location=device)
        model.load_state_dict(state_dict)
        model.to(device)
        model.eval()
        print("   ✅ ConvNeXt Fallback Ready.")
        return model
    except Exception as e:
        print(f"   ❌ Image Model Error: {e}")
        return None

# --- INITIALIZATION ---
print("\n⚙️ --- LOADING HEAVY MODELS INTO MEMORY ---")

# 1. Load Text Models (7D Fallback)
models["text"].append(load_text_model(TEXT_MODEL_1_ID))
models["text"].append(load_text_model(TEXT_MODEL_2_ID))
models["text"].append(load_text_model(TEXT_MODEL_3_ID))

try:
    models["ensemble"] = joblib.load(ENSEMBLE_PATH)
    print("   ✅ English 7D Ensemble Loaded.")
except Exception:
    models["ensemble"] = None

# 2. Load Marathi Models
marathi_components = {}
try:
    marathi_components['ensemble'] = joblib.load(MARATHI_ENSEMBLE_PATH) if os.path.exists(MARATHI_ENSEMBLE_PATH) else None
    marathi_components['tok_maha'] = AutoTokenizer.from_pretrained(MARATHI_MAHABERT_ID, trust_remote_code=True)
    marathi_components['mod_maha'] = AutoModelForSequenceClassification.from_pretrained(MARATHI_MAHABERT_ID, trust_remote_code=True, low_cpu_mem_usage=True).to(device)
    marathi_components['mod_maha'].eval()
    marathi_components['tok_indic'] = AutoTokenizer.from_pretrained(MARATHI_INDICBERT_ID, trust_remote_code=True)
    marathi_components['mod_indic'] = AutoModelForSequenceClassification.from_pretrained(MARATHI_INDICBERT_ID, trust_remote_code=True, low_cpu_mem_usage=True).to(device)
    marathi_components['mod_indic'].eval()
    models["marathi"] = marathi_components
    print("   ✅ Marathi System Ready!")
except Exception:
    models["marathi"] = None

# 3. Load Code Model
try:
    models["code"] = {
        "tokenizer": AutoTokenizer.from_pretrained(CODE_MODEL_ID, trust_remote_code=True),
        "model": AutoModelForSequenceClassification.from_pretrained(CODE_MODEL_ID, trust_remote_code=True, low_cpu_mem_usage=True).to(device)
    }
    models["code"]["model"].eval()
    print("   ✅ Code Model Ready.")
except Exception:
    pass

# 4. Load Image Fallback Model
models["image"] = load_image_model_from_hub(IMAGE_REPO_ID)

gc.collect()
if device == "cuda": torch.cuda.empty_cache()


# ==========================================================
# 3. HELPER FUNCTIONS
# ==========================================================

def get_stylometric_features(text):
    if not text: return [0,0,0,0]
    prob = [float(text.count(c)) / len(text) for c in dict.fromkeys(list(text))]
    entropy = - sum([p * math.log(p) / math.log(2.0) for p in prob])
    sentences = text.replace('!', '.').replace('?', '.').split('.')
    lengths = [len(s.split()) for s in sentences if len(s.split()) > 0]
    burstiness = np.std(lengths) if lengths else 0
    words = text.lower().split()
    ttr = len(set(words)) / len(words) if words else 0
    if len(words) < 3: ngram_ratio = 0
    else:
        ngrams = list(zip(*[words[i:] for i in range(3)]))
        counts = Counter(ngrams)
        repeated = sum(1 for count in counts.values() if count > 1)
        ngram_ratio = repeated / len(ngrams)
    return [entropy, burstiness, ttr, ngram_ratio]

def get_prob(text, tokenizer, model):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = F.softmax(outputs.logits, dim=-1)
    return probs[0][1].item()

def get_marathi_features(text):
    if not text: return [0.0] * 8
    sentences = nltk.sent_tokenize(text)
    words = nltk.word_tokenize(text)
    num_words = len(words)
    if num_words == 0: return [0.0] * 8
    num_sentences = max(len(sentences), 1)

    avg_sentence_len = num_words / num_sentences
    sentence_lengths = [len(nltk.word_tokenize(s)) for s in sentences]
    std_sentence_len = np.std(sentence_lengths) if len(sentence_lengths) > 1 else 0.0
    ttr = len(set(words)) / num_words
    
    counts = Counter(text)
    entropy = -sum((c/len(text)) * math.log2(c/len(text)) for c in counts.values())
    
    word_counts = Counter(words)
    hapax_ratio = sum(1 for c in word_counts.values() if c == 1) / num_words
    avg_word_len = len(text) / num_words
    burstiness = (std_sentence_len / avg_sentence_len) if avg_sentence_len > 0 else 0.0
    
    return [avg_sentence_len, std_sentence_len, ttr, entropy, hapax_ratio, avg_word_len, burstiness, (entropy * avg_sentence_len)]

def get_forensics(img_pil):
    img_np = np.array(img_pil)
    if len(img_np.shape) == 2: img_np = cv2.cvtColor(img_np, cv2.COLOR_GRAY2RGB)
    elif img_np.shape[2] == 4: img_np = cv2.cvtColor(img_np, cv2.COLOR_RGBA2RGB)
         
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    dft = np.fft.fft2(gray)
    dft_shift = np.fft.fftshift(dft)
    magnitude_spectrum = np.log(np.abs(dft_shift) + 1)
    
    return {
        "spectral_artifacts": round(float(np.mean(magnitude_spectrum)), 3),
        "perplexity": round(float(shannon_entropy(gray)), 3),
        "burstiness": round(float(np.std(cv2.Canny(gray, 100, 200))), 3)
    }

def get_image_transforms():
    return transforms.Compose([
        transforms.Resize((384, 384)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

# --- MASTER IMAGE PIPELINE (API + LOCAL FALLBACK) ---
def analyze_image_pipeline(pil_img, filename="image.jpg"):
    global current_sightengine_idx
    keys_tried = 0
    total_keys = len(SIGHTENGINE_API_KEYS)

    # 1. Prepare for API
    img_bytes = io.BytesIO()
    pil_img.save(img_bytes, format='JPEG')

    # 2. Try Sightengine API Rotation
    while keys_tried < total_keys:
        creds = SIGHTENGINE_API_KEYS[current_sightengine_idx]
        try:
            img_bytes.seek(0)
            files = {'media': (filename, img_bytes, 'image/jpeg')}
            params = {'models': 'genai', 'api_user': creds['user'], 'api_secret': creds['secret']}
            
            response = requests.post(SIGHTENGINE_URL, files=files, data=params, timeout=20)
            
            if 200 <= response.status_code < 300:
                result = response.json()
                if result.get("status") == "success":
                    ai_score = result.get('type', {}).get('ai_generated', 0.0)
                    return float(ai_score), f"sightengine_api (Key {current_sightengine_idx + 1})"
                else:
                    current_sightengine_idx = (current_sightengine_idx + 1) % total_keys
                    keys_tried += 1
            else:
                current_sightengine_idx = (current_sightengine_idx + 1) % total_keys
                keys_tried += 1

        except Exception:
            current_sightengine_idx = (current_sightengine_idx + 1) % total_keys
            keys_tried += 1

    # 3. Trigger Local ConvNeXt Fallback
    print(f"❌ Sightengine APIs exhausted/failed for {filename}. Triggering ConvNeXt Fallback.")
    if models.get("image") is not None:
        try:
            transform = get_image_transforms()
            img_t = transform(pil_img).unsqueeze(0).to(device)
            with torch.no_grad():
                logits = models["image"](img_t)
                probs = F.softmax(logits, dim=1)
                # Standard ConvNeXt Index 0 is often AI. Adjust if your model expects Index 1.
                fallback_score = probs[0][0].item() 
            return float(fallback_score), "local_convnext_fallback"
        except Exception as e:
            print(f"⚠️ ConvNeXt Fallback Failed: {e}")
            
    return 0.0, "error_fallback"

# ==========================================================
# 4. API ENDPOINTS
# ==========================================================

class DetectionRequest(BaseModel):
    text: str

@app.post("/analyze")
def analyze_text(request: DetectionRequest):
    global current_sapling_idx
    user_text = request.text
    if len(user_text.strip()) < 5: return {"ai_score": 0, "label": "Too Short", "stats": {}}
    if "def " in user_text and ("return" in user_text or "class" in user_text):
        return {"ai_score": 0.0, "label": "Use /analyze_code endpoint", "stats": {}}

    stats = get_stylometric_features(user_text)
    final_prob = None
    used_fallback = False
    used_source = ""

    api_success = False
    keys_tried = 0
    total_keys = len(SAPLING_API_KEYS)

    while keys_tried < total_keys:
        test_key = SAPLING_API_KEYS[current_sapling_idx]
        try:
            response = requests.post(SAPLING_URL, json={"key": test_key, "text": user_text}, timeout=20)
            if 200 <= response.status_code < 300:
                final_prob = response.json().get("score", 0.0)
                used_source = f"sapling_api (Key {current_sapling_idx + 1})"
                api_success = True
                break 
            else:
                current_sapling_idx = (current_sapling_idx + 1) % total_keys
                keys_tried += 1
        except Exception:
            current_sapling_idx = (current_sapling_idx + 1) % total_keys
            keys_tried += 1

    if not api_success:
        print("❌ Sapling APIs exhausted. Triggering Heavy 7D Fallback.")
        used_fallback = True

    if used_fallback or final_prob is None:
        dl_probs = []
        for entry in models["text"]:
            if entry:
                try:
                    inputs = entry["tokenizer"](user_text, return_tensors="pt", truncation=True, max_length=512).to(device)
                    with torch.no_grad():
                        outputs = entry["model"](**inputs)
                        dl_probs.append(F.softmax(outputs.logits, dim=-1)[0][1].item())
                except Exception:
                    dl_probs.append(0.5) 
            else:
                dl_probs.append(0.5) 

        final_prob = dl_probs[0] 
        if models["ensemble"]:
            try:
                final_prob = float(models["ensemble"].predict_proba(np.array([dl_probs + stats]))[0][1])
            except: pass
                
        used_source = "local_7D_ensemble"
        gc.collect()
        if device == "cuda": torch.cuda.empty_cache()

    return {
        "ai_score": round(float(final_prob), 4),
        "label": "🤖 AI GENERATED" if final_prob > 0.5 else "👤 HUMAN WRITTEN",
        "detailed_scores": {used_source: round(final_prob, 4)},
        "stats": {
            "entropy": round(stats[0], 2), "burstiness": round(stats[1], 2),
            "ttr": round(stats[2], 2), "ngram_ratio": round(stats[3], 2)
        }
    }

@app.post("/analyze_marathi")
async def analyze_marathi(request: DetectionRequest):
    text = request.text.strip()
    if not text: return {"ai_score": 0, "label": "Empty Text", "stats": {}}
    m_sys = models["marathi"]
    if not m_sys: raise HTTPException(status_code=503, detail="Marathi models failed to load.")

    try:
        prob_maha = get_prob(text, m_sys['tok_maha'], m_sys['mod_maha'])
        prob_indic = get_prob(text, m_sys['tok_indic'], m_sys['mod_indic'])
        stats_features = get_marathi_features(text)
        
        confidence = 0.0
        if m_sys['ensemble']:
            try:
                confidence = float(m_sys['ensemble'].predict_proba(np.array([[prob_maha, prob_indic] + stats_features]))[0][1])
            except:
                confidence = 1.0 if m_sys['ensemble'].predict(np.array([[prob_maha, prob_indic] + stats_features]))[0] == 1 else 0.0
        else:
            confidence = (prob_maha + prob_indic) / 2.0

        return {
            "ai_score": round(confidence, 4),
            "label": "🤖 AI GENERATED" if confidence > 0.5 else "👤 HUMAN WRITTEN",
            "stats": {"avg_sent_len": round(stats_features[0], 2), "entropy": round(stats_features[3], 2)}
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/analyze_code")
async def analyze_code(request: DetectionRequest):
    if not models["code"]: raise HTTPException(status_code=503, detail="Code model not loaded.")
    try:
        inputs = models["code"]["tokenizer"](request.text, return_tensors="pt", truncation=True, max_length=512).to(device)
        with torch.no_grad():
            ai_prob = F.softmax(models["code"]["model"](**inputs).logits, dim=-1)[0][1].item()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    stats = get_stylometric_features(request.text)
    return {
        "ai_score": round(float(ai_prob), 4),
        "label": "🤖 AI CODE" if ai_prob > 0.5 else "👤 HUMAN CODE",
        "stats": {"entropy": round(stats[0], 2), "burstiness": round(stats[1], 2)}
    }

@app.post("/analyze_image")
def analyze_image(file: UploadFile = File(...)):
    try:
        pil_img = Image.open(io.BytesIO(file.file.read())).convert('RGB')
        forensics = get_forensics(pil_img)
        
        # MAGIC HAPPENS HERE: Auto-Routes to API -> Falls back to ConvNeXt
        ai_score, used_source = analyze_image_pipeline(pil_img, filename=file.filename)
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        file.file.close()

    return {
        "ai_score": round(float(ai_score), 4),
        "label": "AI Generated" if ai_score > 0.5 else "Real / Human",
        "detailed_scores": {used_source: round(float(ai_score), 4)},
        "forensics": forensics
    }

@app.post("/analyze_video")
def analyze_video(file: UploadFile = File(...), num_samples: int = 10):
    unique_name = f"temp_vid_{uuid.uuid4()}.mp4"
    try:
        with open(unique_name, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
            
        cap = cv2.VideoCapture(unique_name)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if total_frames < 1: raise ValueError("Empty video")
            
        indices = np.linspace(0, total_frames-1, num=min(num_samples, total_frames), dtype=int)
        scores = []
        
        for i in indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
            ret, frame = cap.read()
            if not ret: continue
            
            pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            
            # API -> Fallback
            score, _ = analyze_image_pipeline(pil_img, filename=f"frame_{i}.jpg")
            scores.append(score)
        
        cap.release()
        if not scores: return {"ai_score": 0, "label": "Error"}
        
        avg_score = sum(scores) / len(scores)
        return {
            "ai_score": round(avg_score, 4),
            "label": "AI Video" if avg_score > 0.5 else "Real Video",
            "frames_analyzed": len(scores)
        }
        
    except Exception as e:
        return {"error": str(e)}
    finally:
        if os.path.exists(unique_name): os.remove(unique_name)
        file.file.close()

@app.post("/analyze_document")
def analyze_document(file: UploadFile = File(...)):
    global current_sapling_idx
    try:
        content = file.file.read()
        file_bytes = io.BytesIO(content)

        try:
            doc = Document(file_bytes)
            full_text = "\n".join([para.text for para in doc.paragraphs])
        except: full_text = ""

        extracted_images = []
        try:
            file_bytes.seek(0)
            with zipfile.ZipFile(file_bytes) as z:
                image_files = [f for f in z.namelist() if f.startswith('word/media/')]
                for img_path in image_files:
                    with z.open(img_path) as img_file:
                        try:
                            extracted_images.append({
                                "filename": os.path.basename(img_path),
                                "image_obj": Image.open(io.BytesIO(img_file.read())).convert('RGB')
                            })
                        except: continue
        except: pass

        text_result = None
        if len(full_text.split()) > 5:
            api_success = False
            keys_tried = 0
            t_prob = 0.0

            while keys_tried < len(SAPLING_API_KEYS):
                try:
                    response = requests.post(
                        SAPLING_URL,
                        json={"key": SAPLING_API_KEYS[current_sapling_idx], "text": full_text[:2000]},
                        timeout=20
                    )
                    if 200 <= response.status_code < 300:
                        t_prob = response.json().get("score", 0.0)
                        api_success = True
                        break 
                    else:
                        current_sapling_idx = (current_sapling_idx + 1) % len(SAPLING_API_KEYS)
                        keys_tried += 1
                except Exception:
                    current_sapling_idx = (current_sapling_idx + 1) % len(SAPLING_API_KEYS)
                    keys_tried += 1

            if not api_success:
                stats = get_stylometric_features(full_text)
                dl_probs = []
                for entry in models["text"]:
                    if entry:
                        try:
                            inputs = entry["tokenizer"](full_text[:2000], return_tensors="pt", truncation=True, max_length=512).to(device)
                            with torch.no_grad():
                                dl_probs.append(F.softmax(entry["model"](**inputs).logits, dim=-1)[0][1].item())
                        except: dl_probs.append(0.5)
                    else: dl_probs.append(0.5)
                
                t_prob = dl_probs[0]
                if models["ensemble"]:
                    try: t_prob = float(models["ensemble"].predict_proba(np.array([dl_probs + stats]))[0][1])
                    except: pass
                        
            text_result = {"ai_score": t_prob, "preview": full_text[:100]}

        image_results = []
        for item in extracted_images:
            # API -> Fallback applied to documents too
            img_ai_score, _ = analyze_image_pipeline(item['image_obj'], filename=item['filename'])
            image_results.append({
                "filename": item['filename'],
                "ai_score": round(float(img_ai_score), 4),
                "label": "AI Generated" if img_ai_score > 0.5 else "Real / Human",
                "forensics": get_forensics(item['image_obj'])
            })

        scores = [text_result['ai_score']] if text_result else []
        scores.extend([img['ai_score'] for img in image_results])
        
        final_score = (sum(scores) / len(scores) + max(scores)) / 2 if scores else 0.0

        return {
            "type": "document_report",
            "ai_score": round(final_score, 4),
            "label": "AI DETECTED" if final_score > 0.5 else "HUMAN WRITTEN",
            "text_analysis": text_result,
            "image_analysis": image_results,
            "summary": f"Analyzed {len(full_text.split())} words and {len(image_results)} images."
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Document parsing failed: {str(e)}")
    finally:
        file.file.close()
        
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)