import os
import re
import string
import emoji
import torch
import torch.nn as nn
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

# -----------------------------
# 1️⃣ Set writable cache folder before importing Transformers
# -----------------------------
os.environ["HF_HOME"] = "/tmp/hf_cache"
os.makedirs("/tmp/hf_cache", exist_ok=True)

# -----------------------------
# 2️⃣ Define model class
# -----------------------------
class EnhancedBERT(nn.Module):
    def __init__(self, hidden_dropout=0.3, bert_model=None):
        super(EnhancedBERT, self).__init__()
        self.bert = bert_model
        self.dropout = nn.Dropout(hidden_dropout)
        self.fc = nn.Linear(self.bert.config.hidden_size, 2)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        pooled_output = self.dropout(pooled_output)
        logits = self.fc(pooled_output)
        return logits

# -----------------------------
# 3️⃣ FastAPI setup
# -----------------------------
app = FastAPI(title="Hate Speech Detection API")

class TextInput(BaseModel):
    text: str

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = None
tokenizer = None
model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'Enhanced_BERT_Text.pt')

# -----------------------------
# 4️⃣ Load model & tokenizer at startup
# -----------------------------
@app.on_event("startup")
def load_model():
    global model, tokenizer

    from transformers import BertModel, BertTokenizerFast

    cache_dir = "/tmp/hf_cache"

    # Load tokenizer
    tokenizer = BertTokenizerFast.from_pretrained(
        "bert-base-multilingual-cased",
        cache_dir=cache_dir,
        force_download=False
    )

    # Load pre-trained BERT
    bert_model = BertModel.from_pretrained(
        "bert-base-multilingual-cased",
        cache_dir=cache_dir,
        force_download=False
    )

    # Wrap in EnhancedBERT
    model = EnhancedBERT(bert_model=bert_model).to(device)

    if not os.path.exists(model_path):
        raise FileNotFoundError(f"[ERROR] Model file not found: {model_path}")


    
    state_dict = torch.load(model_path, map_location=device)
    # Fix for DataParallel keys
    new_state_dict = {}
    for k, v in state_dict.items():
        if k.startswith("module."):
            new_key = k[len("module."):]  # remove "module." prefix
        else:
            new_key = k
        new_state_dict[new_key] = v
    
    model.load_state_dict(new_state_dict)

    model.eval()
    print("[INFO] Model loaded successfully")

# -----------------------------
# 5️⃣ Preprocessing functions
# -----------------------------
# def clean_text_round1(text):
#     text = text.lower()
#     text = re.sub(r'http\S+|www\S+|https\S+|http\s+', '', text)
#     text = re.sub(r'<.*?>', '', text)
#     text = re.sub(r'\s+', ' ', text)
#     text = re.sub(r'\n', ' ', text)
#     text = re.sub(r'[\'\"“”]', '', text)
#     text = re.sub(r'&', '', text)
#     text = re.sub(r'@', ' ', text)
#     text = re.sub(r'#', ' ', text)
#     text = re.sub(r'\bRT\b', 'retweet', text)
#     text = re.sub(r'\brt\b', '', text)

#     punctuation = string.punctuation.replace('[', '').replace(']', '').replace('(', '').replace(')', '')
#     text = re.sub(r'[%s]' % re.escape(punctuation), '', text)

#     abbreviations = {
#         r'\bthx\b': 'thanks',
#         r'\bpls\b': 'please',
#         r'\bplz\b': 'please',
#         r'\bbtw\b': 'by the way',
#         r'\bomg\b': 'oh my god',
#         r'\bidk\b': 'i dont know',
#         r'\bstfu\b': 'shut the fuck up',
#         r'\blmao+\b': 'laughing my ass off',
#         r'\bidc\b': 'i dont care',
#         r'\bwtf\b': 'what the fuck',
#         r'\baf\b': 'as fuck',
#         r'\blol\b': 'laugh out loud',
#         r'\bgtfo\b': 'get the fuck out',
#         r'\bnvm\b': 'never mind',
#         r'\brip\b': 'rest in peace'
#     }
#     for abbr, full in abbreviations.items():
#         text = re.sub(abbr, full, text)

#     return text.strip()

# def clean_demojized_text(text):
#     return emoji.demojize(text).replace(":", " ").strip()

def preprocess_text(text, max_len=75):
    encoding = tokenizer(
        text,
        add_special_tokens=True,
        max_length=max_len,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    )
    return encoding

def predict_text_input(text):
    encoding = preprocess_text(text)
    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)

    with torch.no_grad():
        logits = model(input_ids, attention_mask)
        probabilities = torch.softmax(logits, dim=1).cpu().numpy()[0]
        label_idx = probabilities.argmax()
        label = "Hate Speech" if label_idx == 1 else "Non-Hate Speech"
        probability = probabilities[label_idx]

    return {
        'prediction': label,
        'probability': float(probability)
    }

# -----------------------------
# 6️⃣ Routes
# -----------------------------
@app.get("/")
async def root():
    return {"message": "API is running. Use POST /predict"}

@app.post("/predict")
@app.post("/predict/")
async def predict_route(input_data: TextInput):
    result = predict_text_input(input_data.text)
    return {
        "label": result['prediction'],
        "probability": round(result['probability'], 4)
    }

# -----------------------------
# 7️⃣ Run server
# -----------------------------
if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)