import os
import shutil
import torch
import torch.nn as nn
import torch.nn.functional as F
import logging
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer, AutoModel

# Configs
REPO_ID = "can-org/Nepali-AI-VS-HUMAN"
BASE_DIR = "./np_text_model"
TOKENIZER_DIR = os.path.join(BASE_DIR, "classifier")  # <- update this to match your uploaded folder
WEIGHTS_PATH = os.path.join(BASE_DIR, "model_95_acc.pth")  # <- change to match actual uploaded weight
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Define model class
class XLMRClassifier(nn.Module):
    def __init__(self):
        super(XLMRClassifier, self).__init__()
        self.bert = AutoModel.from_pretrained("xlm-roberta-base")
        self.classifier = 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)
        cls_output = outputs.last_hidden_state[:, 0, :]
        return self.classifier(cls_output)

# Globals for caching
_model = None
_tokenizer = None

def download_model_repo():
    if os.path.exists(BASE_DIR) and os.path.isdir(BASE_DIR):
        logging.info("Model already downloaded.")
        return
    snapshot_path = snapshot_download(repo_id=REPO_ID)
    os.makedirs(BASE_DIR, exist_ok=True)
    shutil.copytree(snapshot_path, BASE_DIR, dirs_exist_ok=True)

def load_model():
    download_model_repo()
    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
    model = XLMRClassifier().to(device)
    model.load_state_dict(torch.load(WEIGHTS_PATH, map_location=device))
    model.eval()
    return model, tokenizer

def get_model_tokenizer():
    global _model, _tokenizer
    if _model is None or _tokenizer is None:
        _model, _tokenizer = load_model()
    return _model, _tokenizer