"""
RuBERT Toxicity Detection API for HuggingFace Spaces
FastAPI + Gradio solution for REST API compatibility
"""

import os
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import json
import time
from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel
import uvicorn
from typing import List, Optional
from collections import defaultdict

# Security configuration
API_KEY = os.getenv("API_KEY", "")  # Set via HuggingFace Spaces secrets
MAX_TEXT_LENGTH = 2000  # Max characters per request
RATE_LIMIT_WINDOW = 60  # Seconds
RATE_LIMIT_MAX_REQUESTS = 30  # Max requests per window

# Simple in-memory rate limiter
rate_limit_store = defaultdict(list)

def check_rate_limit(client_id: str) -> bool:
    """Check if client has exceeded rate limit"""
    now = time.time()
    rate_limit_store[client_id] = [
        t for t in rate_limit_store[client_id]
        if now - t < RATE_LIMIT_WINDOW
    ]
    if len(rate_limit_store[client_id]) >= RATE_LIMIT_MAX_REQUESTS:
        return False
    rate_limit_store[client_id].append(now)
    return True

def verify_api_key(authorization: Optional[str]) -> bool:
    """Verify API key if configured"""
    if not API_KEY:
        return True
    if not authorization:
        return False
    key = authorization.replace("Bearer ", "").strip()
    return key == API_KEY

# Load model and tokenizer
print("Loading RuBERT model...")
model_name = "cointegrated/rubert-tiny-toxicity"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval()

# Labels from the model
LABELS = ['non-toxic', 'insult', 'obscenity', 'threat', 'dangerous']

def classify_text(text):
    """
    Classify text for toxicity using RuBERT
    Returns dictionary with classification scores
    """
    start_time = time.time()

    # Tokenize and predict
    with torch.no_grad():
        inputs = tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            max_length=512
        )
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)

    # Convert to dictionary
    scores = {}
    for i, label in enumerate(LABELS):
        scores[label] = float(probs[0][i])

    # Calculate toxicity based on model's formula
    # toxicity = 1 - (non_toxic * (1 - dangerous))
    toxicity_score = 1 - (scores['non-toxic'] * (1 - scores['dangerous']))

    # Determine if toxic
    is_toxic = (
        toxicity_score > 0.5 or
        scores['insult'] > 0.6 or
        scores['obscenity'] > 0.6 or
        scores['threat'] > 0.6
    )

    # Get toxic labels
    toxic_labels = [
        label for label in LABELS[1:]
        if scores[label] > 0.5
    ]

    # Format response compatible with the bot
    result = {
        "toxic": is_toxic,
        "confidence": toxicity_score if is_toxic else 1 - toxicity_score,
        "scores": scores,
        "labels": toxic_labels,
        "processing_time_ms": (time.time() - start_time) * 1000
    }

    return result

# Create FastAPI app
app = FastAPI(title="RuBERT Toxicity Detection API")

# Request model for API
class PredictRequest(BaseModel):
    data: List[str]

# API endpoint for bot integration
@app.post("/api/predict")
async def api_predict(
    request: PredictRequest,
    authorization: Optional[str] = Header(None),
    x_forwarded_for: Optional[str] = Header(None, alias="X-Forwarded-For")
):
    """
    REST API endpoint compatible with existing bot integration
    Expects: {"data": ["text to analyze"]}
    Returns: {"data": [json_string_with_results]}
    """
    # Verify API key if configured
    if not verify_api_key(authorization):
        raise HTTPException(status_code=401, detail="Invalid or missing API key")

    # Rate limiting
    client_id = x_forwarded_for or "default"
    if not check_rate_limit(client_id):
        raise HTTPException(status_code=429, detail="Rate limit exceeded")

    if not request.data or len(request.data) == 0:
        return {"error": "No text provided"}

    text = request.data[0]

    # Input size limit
    if len(text) > MAX_TEXT_LENGTH:
        text = text[:MAX_TEXT_LENGTH]

    result = classify_text(text)

    # Return in Gradio-compatible format
    return {
        "data": [json.dumps(result, ensure_ascii=False)]
    }

# Gradio interface wrapper
def gradio_classify(text):
    """Wrapper for Gradio that returns JSON string"""
    result = classify_text(text)
    return json.dumps(result, ensure_ascii=False, indent=2)

# Create Gradio interface
iface = gr.Interface(
    fn=gradio_classify,
    inputs=gr.Textbox(
        lines=3,
        placeholder="Введите текст для проверки на токсичность / Enter text to check for toxicity",
        label="Text"
    ),
    outputs=gr.JSON(label="Classification Results"),
    title="🤖 RuBERT Toxicity Detection",
    description="""
    Russian toxicity detection using cointegrated/rubert-tiny-toxicity model.

    Detects: insults, obscenity, threats, and dangerous content in Russian text.

    API Usage:
    ```python
    import requests
    response = requests.post(
        "https://tsmtgkkehawwcnahpmyl-rubert-toxicity-detector.hf.space/api/predict",
        json={"data": ["your text here"]}
    )
    result = response.json()
    ```
    """,
    examples=[
        ["Привет, как дела?"],
        ["Ты меня затрахал"],
        ["Какой же ты мудак"],
        ["Спасибо за помощь"],
    ]
)

# Mount Gradio app to FastAPI
app = gr.mount_gradio_app(app, iface, path="/")

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