from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
import joblib
import json
import pandas as pd
import numpy as np
import re
from transformers import AutoTokenizer, AutoModelForSequenceClassification, BertConfig
from huggingface_hub import hf_hub_download
import os
from safetensors.torch import load_file
import uvicorn

# Initialize FastAPI app
app = FastAPI(title="Ticket Field Predictor API")

# Get Hugging Face token (set via Space Secrets as HF_TOKEN)
HUGGINGFACE_TOKEN = os.getenv("HF_TOKEN")

# Model Repository on Hugging Face
REPO_ID = "AtaishNehra/ticket_field_predictor_model"

# Define features
FEATURES = [
    "Company", "Contract", "Priority", "Issue Type",
    "Service Level Agreement", "Sub-Issue Type", "Queue"
]

# Load label encoders
encoder_path = hf_hub_download(repo_id=REPO_ID, filename="label_encoders.joblib", token=HUGGINGFACE_TOKEN)
label_encoders = joblib.load(encoder_path)

# Load tokenizers
tokenizers = {feature: AutoTokenizer.from_pretrained("bert-base-uncased") for feature in FEATURES}

# Function to load models
def load_model(feature):
    config_filename = f"{feature}/config.json"
    model_filename = f"{feature}/model.safetensors"

    # Download and load the configuration
    config_path = hf_hub_download(repo_id=REPO_ID, filename=config_filename, token=HUGGINGFACE_TOKEN)
    with open(config_path, "r") as f:
        config = json.load(f)

    # Download and load the model state dict
    model_path = hf_hub_download(repo_id=REPO_ID, filename=model_filename, token=HUGGINGFACE_TOKEN)
    state_dict = load_file(model_path, device="cpu")

    # Ensure configuration has a model_type
    config.setdefault("model_type", "bert")
    model_config = BertConfig.from_dict(config)
    model = AutoModelForSequenceClassification.from_config(model_config)

    # Load state dict without strict checking
    model.load_state_dict(state_dict, strict=False)
    model.eval()
    return model

# Load models for each feature
models = {feature: load_model(feature) for feature in FEATURES}

# Input model for API request
class TicketInput(BaseModel):
    title: str
    description: str

# Prediction function
@torch.no_grad()
def predict(title: str, description: str):
    inputs = f"{title} {description}"
    predictions = {}

    for feature in FEATURES:
        tokenizer = tokenizers[feature]
        model = models[feature]
        label_encoder = label_encoders[feature]

        tokens = tokenizer(inputs, return_tensors="pt", padding=True, truncation=True)
        output = model(**tokens).logits
        pred_idx = torch.argmax(output, dim=1).item()

        # Decode prediction
        predictions[feature] = label_encoder.inverse_transform([pred_idx])[0]

    return predictions

@app.post("/predict")
async def api_predict(ticket: TicketInput):
    """API endpoint for external HTTP requests"""
    try:
        predictions = predict(ticket.title, ticket.description)
        return {"predictions": predictions}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

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