import uvicorn
from fastapi import FastAPI, HTTPException, Security, Depends
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import tensorflow as tf
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
import os
import jwt

# Initialize FastAPI app
app = FastAPI(title="AI vs Human Detector API")

# CORS Middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

# Security
security = HTTPBearer()
SECRET_KEY = "mysecretkey"  # In production, use environment variable

# Global variables for model and tokenizer
model = None
tokenizer = None
MODEL_PATH = "saved_human_ai_model"

class PredictionRequest(BaseModel):
    text: str

class PredictionResponse(BaseModel):
    label: str
    confidence: float
    probabilities: dict

def verify_token(credentials: HTTPAuthorizationCredentials = Security(security)):
    token = credentials.credentials
    try:
        # Just verify the signature using the SECRET_KEY
        # We don't need to parse/use the payload for user verification
        jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
    except jwt.ExpiredSignatureError:
        raise HTTPException(status_code=401, detail="Token has expired")
    except jwt.InvalidTokenError:
        raise HTTPException(status_code=401, detail="Invalid token")

@app.on_event("startup")
async def load_model():
    global model, tokenizer
    try:
        print(f"Loading model from {MODEL_PATH}...")
        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
        model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
        print("Model and Tokenizer loaded successfully!")
    except Exception as e:
        print(f"Error loading model: {e}")
        raise RuntimeError(f"Could not load model: {e}")

@app.post("/predict", response_model=PredictionResponse, dependencies=[Depends(verify_token)])
async def predict(request: PredictionRequest):
    if not model or not tokenizer:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    try:
        # Tokenize input
        inputs = tokenizer(
            request.text,
            return_tensors="tf",
            padding=True,
            truncation=True,
            max_length=512
        )
        
        # Inference
        outputs = model(inputs)
        logits = outputs.logits
        
        # Softmax
        probabilities = tf.nn.softmax(logits, axis=-1).numpy()[0]
        
        # Get prediction
        predicted_class_id = np.argmax(probabilities)
        confidence = float(probabilities[predicted_class_id])
        
        # Map labels (Assuming 0=Human, 1=AI based on notebook)
        labels_map = {0: "Human", 1: "AI"}
        predicted_label = labels_map.get(predicted_class_id, "Unknown")
        
        return PredictionResponse(
            label=predicted_label,
            confidence=confidence,
            probabilities={
                "Human": float(probabilities[0]),
                "AI": float(probabilities[1])
            }
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# if __name__ == "__main__":
#     uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)