import os
import logging
import random
import pickle
import torch
from transformers import BertForSequenceClassification, BertTokenizerFast
from huggingface_hub import hf_hub_download
import gradio as gr

# Configuration
MODEL_NAME = "vinapatri/intent-classification-jkn-kis"
CONFIDENCE_THRESHOLD = 0.6  # Adjust this based on your model's performance

# --- Model Functions ---
def load_model_components():
    model = BertForSequenceClassification.from_pretrained(MODEL_NAME)
    tokenizer = BertTokenizerFast.from_pretrained(MODEL_NAME)
    
    label_encoder_path = hf_hub_download(repo_id=MODEL_NAME, filename="label_encoder.pkl")
    responses_path = hf_hub_download(repo_id=MODEL_NAME, filename="tag_to_responses.pkl")
    
    with open(label_encoder_path, "rb") as f:
        le = pickle.load(f)
    
    with open(responses_path, "rb") as f:
        tag_to_responses = pickle.load(f)
    
    # Ensure we have an 'unknown' response category
    if 'unknown' not in tag_to_responses:
        tag_to_responses['unknown'] = [
            "Maaf, saya tidak mengerti pertanyaan Anda",
            "Saya belum bisa menjawab pertanyaan tersebut"
        ]
    
    return model, tokenizer, le, tag_to_responses

model, tokenizer, le, tag_to_responses = load_model_components()

def predict_intent(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Get confidence scores
    probabilities = torch.softmax(outputs.logits, dim=1)[0]
    confidence, predicted_class_id = torch.max(probabilities, 0)
    confidence = confidence.item()
    
    # Return 'unknown' if confidence is below threshold
    if confidence < CONFIDENCE_THRESHOLD:
        return "unknown"
    
    tag = le.inverse_transform([predicted_class_id.item()])[0]
    return tag

def get_response(user_input):
    tag = predict_intent(user_input)
    responses = tag_to_responses.get(tag, tag_to_responses['unknown'])
    return random.choice(responses)

# --- Gradio Interface ---
def chat_interface(message, history):
    return get_response(message)

gradio_app = gr.ChatInterface(
    fn=chat_interface,
    title="JKN-KIS Intent Classification",
    description="Bot untuk klasifikasi intent terkait JKN-KIS",
    examples=[
        "Bagaimana cara daftar Mobile JKN?",
        "Cara buat surat rujukan BPJS secara online?",
        "Kenapa OTP Mobile JKN tidak masuk?"
    ]
)

if __name__ == "__main__":
    logging.basicConfig(
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        level=logging.INFO
    )
    gradio_app.launch(server_name="0.0.0.0", server_port=7860)