from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field, validator
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import uvicorn
import logging
import time
import hashlib
from functools import lru_cache
from typing import Optional
import re

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    title="AI Content Classifier API",
    description="Detect whether the given text is human-written, AI-generated, or paraphrased.",
    version="1.0.0"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Configure this properly for production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load the model and tokenizer from Hugging Face Hub
model_name = None
model = None
tokenizer = None

# Try multiple models in order of preference
model_options = [
    "hello-SimpleAI/chatgpt-detector-roberta",
    "roberta-base-openai-detector", 
    "distilbert-base-uncased"
]

for model_option in model_options:
    try:
        logger.info(f"Attempting to load model: {model_option}")
        if model_option == "distilbert-base-uncased":
            # Create a simple 3-class classifier for demo
            model = AutoModelForSequenceClassification.from_pretrained(model_option, num_labels=3, ignore_mismatched_sizes=True)
        else:
            model = AutoModelForSequenceClassification.from_pretrained(model_option)
        
        tokenizer = AutoTokenizer.from_pretrained(model_option)
        model_name = model_option
        logger.info(f"Successfully loaded model: {model_name}")
        break
    except Exception as e:
        logger.warning(f"Failed to load {model_option}: {str(e)}")
        continue

if model is None:
    logger.error("Failed to load any model. Using demo mode.")
    # Create a demo mode flag
    model_name = "demo-mode"

# Configuration constants
MAX_TEXT_LENGTH = 10000
MIN_TEXT_LENGTH = 10
CACHE_SIZE = 1000

# Custom exceptions
class TextTooLongError(HTTPException):
    def __init__(self):
        super().__init__(status_code=400, detail=f"Text exceeds maximum length of {MAX_TEXT_LENGTH} characters")

class TextTooShortError(HTTPException):
    def __init__(self):
        super().__init__(status_code=400, detail=f"Text must be at least {MIN_TEXT_LENGTH} characters long")

# Pydantic models for request/response
class TextRequest(BaseModel):
    text: str = Field(..., min_length=MIN_TEXT_LENGTH, max_length=MAX_TEXT_LENGTH, description="Text to classify")
    
    @validator('text')
    def validate_text_content(cls, v):
        # Remove excessive whitespace
        v = re.sub(r'\s+', ' ', v.strip())
        
        # Check for meaningful content (not just spaces/symbols)
        if len(re.findall(r'[a-zA-Z]', v)) < 5:
            raise ValueError("Text must contain meaningful alphabetic content")
        
        return v

class ClassificationResponse(BaseModel):
    classification: str = Field(..., description="Predicted classification category")
    confidence: float = Field(..., ge=0, le=1, description="Confidence score for the prediction")
    probabilities: dict = Field(..., description="Probability distribution across all categories")
    analysis: dict = Field(..., description="Detailed text analysis metrics")
    suggestions: list = Field(..., description="Improvement suggestions based on classification")
    processing_time: float = Field(..., description="Time taken to process the request in seconds")
    text_hash: str = Field(..., description="Hash of input text for caching purposes")

# Cache for predictions to improve performance
@lru_cache(maxsize=CACHE_SIZE)
def get_cached_prediction(text_hash: str, text: str):
    """Get cached prediction or compute new one"""
    return _classify_text_internal(text)

def create_text_hash(text: str) -> str:
    """Create a hash for the input text for caching"""
    return hashlib.md5(text.encode('utf-8')).hexdigest()[:16]

# Define function for classification
def classify_text(text: str):
    start_time = time.time()
    
    if not text.strip():
        raise HTTPException(status_code=400, detail="Text cannot be empty")
    
    # Create hash for caching
    text_hash = create_text_hash(text)
    
    try:
        # Try to get from cache first
        result = get_cached_prediction(text_hash, text)
        processing_time = time.time() - start_time
        
        classification, confidence, probs_dict, analysis, suggestions = result
        
        # Add processing time to analysis
        analysis["processing_time"] = processing_time
        
        logger.info(f"Text classified as {classification} with confidence {confidence:.4f} in {processing_time:.3f}s")
        
        return classification, confidence, probs_dict, analysis, suggestions, processing_time, text_hash
        
    except Exception as e:
        logger.error(f"Classification error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Classification failed: {str(e)}")

def _demo_classification(text: str):
    """Demo classification when no model is available"""
    import random
    
    # Simple rule-based classification for demo
    words = text.split()
    
    # Check for AI-typical patterns
    ai_indicators = ["furthermore", "therefore", "consequently", "moreover", "in conclusion", "additionally"]
    formal_count = sum(1 for word in words if word.lower() in ai_indicators)
    
    # Calculate synthetic probabilities
    if formal_count > 2:
        # High formality suggests AI
        ai_prob = 0.7 + random.uniform(0, 0.2)
        human_prob = 0.2 + random.uniform(0, 0.1)
        para_prob = 1.0 - ai_prob - human_prob
        predicted_class = 1  # AI-Generated
    elif len(words) > 100 and formal_count > 0:
        # Medium formality suggests paraphrased
        para_prob = 0.5 + random.uniform(0, 0.3)
        ai_prob = 0.3 + random.uniform(0, 0.2)
        human_prob = 1.0 - ai_prob - para_prob
        predicted_class = 2  # Paraphrased
    else:
        # Informal suggests human
        human_prob = 0.6 + random.uniform(0, 0.3)
        ai_prob = 0.2 + random.uniform(0, 0.2)
        para_prob = 1.0 - ai_prob - human_prob
        predicted_class = 0  # Human-Written
    
    # Normalize probabilities
    total = ai_prob + human_prob + para_prob
    probs_dict = {
        "Human-Written": human_prob / total,
        "AI-Generated": ai_prob / total,
        "Paraphrased": para_prob / total
    }
    
    labels = {0: "Human-Written", 1: "AI-Generated", 2: "Paraphrased"}
    classification = labels[predicted_class]
    confidence = max(probs_dict.values())
    
    # Analyze text features
    analysis = analyze_text_features(text, probs_dict, predicted_class)
    analysis["demo_mode"] = True
    analysis["warning"] = "Running in demo mode - results are for demonstration only"
    
    # Generate improvement suggestions
    suggestions = generate_suggestions(predicted_class, probs_dict, text)
    suggestions.insert(0, "⚠️  Demo Mode: Install a proper model for accurate classification")
    
    return classification, confidence, probs_dict, analysis, suggestions

def _classify_text_internal(text: str):
    """Internal classification function without caching logic"""
    
    # Demo mode - when no model could be loaded
    if model_name == "demo-mode":
        return _demo_classification(text)
    
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
    
    logits = outputs.logits
    probabilities = torch.softmax(logits, dim=1)
    predicted_class = torch.argmax(logits, dim=1).item()
    confidence = probabilities[0][predicted_class].item()

    # Determine labels based on model type
    if "openai-detector" in model_name.lower():
        # OpenAI detector: binary classification (Real=0, Fake=1)
        if logits.shape[1] == 2:
            labels = {0: "Human-Written", 1: "AI-Generated"}
            # Add synthetic paraphrased category based on confidence
            probs_dict = {
                "Human-Written": float(probabilities[0][0].item()),
                "AI-Generated": float(probabilities[0][1].item()),
                "Paraphrased": 0.0  # Not supported by this model
            }
        else:
            labels = {0: "Human-Written", 1: "AI-Generated", 2: "Paraphrased"}
            probs_dict = {}
            for i, label in labels.items():
                probs_dict[label] = float(probabilities[0][i].item())
    else:
        # Default 3-class classification or create synthetic classification
        if logits.shape[1] >= 3:
            labels = {0: "Human-Written", 1: "AI-Generated", 2: "Paraphrased"}
        else:
            # For models with different number of classes, create synthetic mapping
            labels = {0: "Human-Written", 1: "AI-Generated"}
        
        probs_dict = {}
        for i, label in labels.items():
            if i < logits.shape[1]:
                probs_dict[label] = float(probabilities[0][i].item())
        
        # Add paraphrased category if not present
        if "Paraphrased" not in probs_dict:
            # Estimate paraphrased probability based on uncertainty
            uncertainty = 1.0 - max(probs_dict.values())
            probs_dict["Paraphrased"] = min(uncertainty, 0.3)  # Cap at 30%
            # Normalize probabilities
            total = sum(probs_dict.values())
            probs_dict = {k: v/total for k, v in probs_dict.items()}
    
    # Map predicted class to label
    classification = labels.get(predicted_class, "Human-Written")
    
    # Analyze text features
    analysis = analyze_text_features(text, probs_dict, predicted_class)
    
    # Generate improvement suggestions
    suggestions = generate_suggestions(predicted_class, probs_dict, text)
    
    return classification, confidence, probs_dict, analysis, suggestions

def analyze_text_features(text: str, probabilities: dict, predicted_class: int):
    """Analyze text features and AI generation probability indicators"""
    words = text.split()
    sentences = [s.strip() for s in text.split('.') if s.strip()]
    
    # Calculate advanced metrics
    avg_word_length = sum(len(word.strip('.,!?;:"()[]')) for word in words) / len(words) if words else 0
    avg_sentence_length = len(words) / len(sentences) if sentences else 0
    
    # Calculate vocabulary diversity (unique words / total words)
    unique_words = set(word.lower().strip('.,!?;:"()[]') for word in words)
    vocab_diversity = len(unique_words) / len(words) if words else 0
    
    # Count punctuation density
    punctuation_count = len(re.findall(r'[.!?;:,]', text))
    punctuation_density = punctuation_count / len(text) if text else 0
    
    analysis = {
        "text_length": len(text),
        "word_count": len(words),
        "sentence_count": len(sentences),
        "avg_word_length": round(avg_word_length, 2),
        "avg_sentence_length": round(avg_sentence_length, 2),
        "vocabulary_diversity": round(vocab_diversity, 4),
        "punctuation_density": round(punctuation_density, 4),
        "ai_indicators": [],
        "human_indicators": [],
        "risk_level": "low",
        "readability_metrics": {
            "complexity_score": round((avg_word_length * avg_sentence_length) / 10, 2),
            "formality_indicators": []
        }
    }
    
    # Common features of AI-generated text
    ai_patterns = [
        "repetitive phrases",
        "overly formal structure", 
        "generic language patterns",
        "lack of personal touch",
        "perfect grammar without variations"
    ]
    
    # Human writing characteristics
    human_patterns = [
        "natural language flow",
        "personal writing style",
        "minor grammatical imperfections",
        "contextual nuances",
        "emotional expressions"
    ]
    
    # Enhanced risk analysis using multiple factors
    ai_prob = probabilities.get("AI-Generated", 0)
    
    # Detect formal/academic language patterns
    formal_indicators = len(re.findall(r'\b(therefore|furthermore|consequently|moreover|additionally|in conclusion|in summary)\b', text.lower()))
    if formal_indicators > 0:
        analysis["readability_metrics"]["formality_indicators"].append(f"Contains {formal_indicators} formal connectors")
    
    # Detect repetitive patterns
    word_freq = {}
    for word in words:
        clean_word = word.lower().strip('.,!?;:"()[]')
        word_freq[clean_word] = word_freq.get(clean_word, 0) + 1
    
    repeated_words = [word for word, freq in word_freq.items() if freq > 3 and len(word) > 3]
    if repeated_words:
        analysis["ai_indicators"].append(f"Repetitive use of words: {', '.join(repeated_words[:3])}")
    
    # Risk level determination with enhanced logic
    risk_factors = 0
    risk_factors += min(ai_prob * 3, 3)  # AI probability weight
    risk_factors += 1 if vocab_diversity < 0.6 else 0  # Low vocabulary diversity
    risk_factors += 1 if avg_sentence_length > 25 else 0  # Very long sentences
    risk_factors += 1 if formal_indicators > 2 else 0  # High formality
    risk_factors += 1 if len(repeated_words) > 2 else 0  # Repetitive language
    
    if risk_factors >= 4:
        analysis["risk_level"] = "high"
        analysis["ai_indicators"].extend([
            f"High AI generation probability ({ai_prob:.3f})",
            "Multiple AI-characteristic patterns detected"
        ])
    elif risk_factors >= 2:
        analysis["risk_level"] = "medium"
        analysis["ai_indicators"].extend([
            f"Moderate AI generation probability ({ai_prob:.3f})",
            "Some AI-characteristic patterns detected"
        ])
    else:
        analysis["risk_level"] = "low"
        analysis["human_indicators"].extend([
            "Natural language flow detected",
            f"Good vocabulary diversity ({vocab_diversity:.3f})"
        ])
    
    return analysis

def generate_suggestions(predicted_class: int, probabilities: dict, text: str):
    """Generate improvement suggestions based on classification results"""
    suggestions = []
    
    ai_prob = probabilities.get("AI-Generated", 0)
    human_prob = probabilities.get("Human-Written", 0)
    paraphrased_prob = probabilities.get("Paraphrased", 0)
    
    if predicted_class == 1:  # AI-Generated
        suggestions.extend([
            "Add more personalized expressions and opinions",
            "Use more natural, informal language style",
            "Include specific examples and personal experiences",
            "Avoid overly perfect grammatical structures",
            "Increase emotional tone and subjective judgments"
        ])
    elif predicted_class == 2:  # Paraphrased
        suggestions.extend([
            "Reorganize article structure to be more original",
            "Add new perspectives and insights",
            "Use more diverse vocabulary and expressions",
            "Include original analysis and conclusions"
        ])
    else:  # Human-Written
        if ai_prob > 0.3:  # Even classified as human-written, but AI probability is high
            suggestions.extend([
                "Maintain current natural writing style",
                "Feel free to express personal opinions more boldly",
                "Continue maintaining natural language fluency"
            ])
    
    # General suggestions
    if len(text.split()) < 50:
        suggestions.append("Text is relatively short; adding more content can improve analysis accuracy")
    
    return suggestions

# API endpoint
@app.post("/detect", response_model=ClassificationResponse)
async def detect_content(request: TextRequest):
    """
    Classify text as human-written, AI-generated, or paraphrased.
    """
    try:
        classification, confidence, probabilities, analysis, suggestions, processing_time, text_hash = classify_text(request.text)
        return ClassificationResponse(
            classification=classification,
            confidence=confidence,
            probabilities=probabilities,
            analysis=analysis,
            suggestions=suggestions,
            processing_time=processing_time,
            text_hash=text_hash
        )
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Unexpected error in detect_content: {str(e)}")
        raise HTTPException(status_code=500, detail="Internal server error")

# Health check endpoint
@app.get("/health")
async def health_check():
    return {
        "status": "healthy",
        "model": model_name,
        "cache_size": get_cached_prediction.cache_info().currsize if hasattr(get_cached_prediction, 'cache_info') else 0,
        "version": "1.0.0"
    }

# Cache statistics endpoint
@app.get("/stats")
async def get_stats():
    cache_info = get_cached_prediction.cache_info() if hasattr(get_cached_prediction, 'cache_info') else None
    return {
        "cache_stats": {
            "hits": cache_info.hits if cache_info else 0,
            "misses": cache_info.misses if cache_info else 0,
            "current_size": cache_info.currsize if cache_info else 0,
            "max_size": cache_info.maxsize if cache_info else CACHE_SIZE
        },
        "model_info": {
            "model_name": model_name,
            "max_text_length": MAX_TEXT_LENGTH,
            "min_text_length": MIN_TEXT_LENGTH
        }
    }

# Clear cache endpoint
@app.post("/admin/clear-cache")
async def clear_cache():
    get_cached_prediction.cache_clear()
    logger.info("Cache cleared")
    return {"message": "Cache cleared successfully"}

# Root endpoint
@app.get("/")
async def root():
    return {
        "message": "AI Content Classifier API is running",
        "version": "1.0.0",
        "endpoints": {
            "classify": "/detect",
            "health": "/health", 
            "stats": "/stats",
            "docs": "/docs"
        }
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)