import os
os.environ["HF_HOME"] = "./app_model/hf_cache"
os.environ["TRANSFORMERS_CACHE"] = "./app_model/hf_cache"
os.environ["HF_DATASETS_CACHE"] = "./app_model/hf_cache"
os.environ["HF_METRICS_CACHE"] = "./app_model/hf_cache"


# Ensure model directory exists
MODEL_DIR = "./app_model"
os.makedirs(MODEL_DIR, exist_ok=True)


from pathlib import Path
import torch
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# ------------------------TRAINING/FINETUNING------------------------

# refer to notebook for the training

# ------------------------INFERENCE------------------------



model_path = snapshot_download(
    repo_id="milanchndr/email-classification-model",
    local_dir=MODEL_DIR,
    local_dir_use_symlinks=False
)


# Load model and tokenizer from the local path
model = AutoModelForSequenceClassification.from_pretrained(model_path, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)


model.eval()  # Set model to evaluation mode

label_map = {0: "Incident", 1: "Request", 2: "Change", 3: "Problem"}

def classify_email(email: str) -> str:
    """Classify an email into a support category using a fine-tuned model.

    Args:
        email (str): The email text to classify.

    Returns:
        str: The predicted category (Incident, Request, Change, or Problem).
    """
    inputs = tokenizer(
        email, padding=True, truncation=True, max_length=512, return_tensors="pt"
    )
    with torch.no_grad():
        outputs = model(**inputs)
        probs = F.softmax(outputs.logits, dim=1)
        pred = torch.argmax(probs, dim=1).item()
    return label_map[pred]