import os
import json
import numpy as np
from tokenizers import Tokenizer
import onnxruntime as ort
from huggingface_hub import hf_hub_download
import gradio as gr

class ONNXInferencePipeline:
    def __init__(self, repo_id):
        # Retrieve the Hugging Face token from the environment variable
        hf_token = os.getenv("HF_TOKEN")
        if hf_token is None:
            raise ValueError("HF_TOKEN environment variable is not set.")

        # Download files from Hugging Face Hub using the token
        self.onnx_path = hf_hub_download(
            repo_id=repo_id,
            filename="moodlens.onnx",
            token=hf_token,  # <-- fixed
        )
        self.tokenizer_path = hf_hub_download(
            repo_id=repo_id,
            filename="train_bpe_tokenizer.json",
            token=hf_token,  # <-- fixed
        )
        self.config_path = hf_hub_download(
            repo_id=repo_id,
            filename="hyperparameters.json",
            token=hf_token,  # <-- fixed
        )

        # Load configuration
        with open(self.config_path) as f:
            self.config = json.load(f)

        # Initialize tokenizer
        self.tokenizer = Tokenizer.from_file(self.tokenizer_path)
        self.max_len = self.config["tokenizer"]["max_len"]

        # Initialize ONNX runtime session
        self.session = ort.InferenceSession(self.onnx_path)
        self.providers = ["CPUExecutionProvider"]  # Use CUDA if available
        if "CUDAExecutionProvider" in ort.get_available_providers():
            self.providers = ["CUDAExecutionProvider"]
        self.session.set_providers(self.providers)

    def preprocess(self, text):
        encoding = self.tokenizer.encode(text)
        ids = encoding.ids[: self.max_len]
        padding = [0] * (self.max_len - len(ids))
        return np.array(ids + padding, dtype=np.int64).reshape(1, -1)

    def predict(self, text):
        # Preprocess
        input_array = self.preprocess(text)

        # Run inference
        results = self.session.run(None, {"input": input_array})

        # Post-process
        logits = results[0]
        probabilities = np.exp(logits) / np.sum(np.exp(logits), axis=1, keepdims=True)
        predicted_class = int(np.argmax(probabilities))

        # Map to labels
        label_mapping = {"neg": "Negative", "pos": "Positive"}
        class_labels = ["neg", "pos"]
        return {
            "label": label_mapping[class_labels[predicted_class]],
            "confidence": float(probabilities[0][predicted_class]),
            "probabilities": probabilities[0].tolist(),
        }

if __name__ == "__main__":
    pipeline = ONNXInferencePipeline(repo_id="iimran/Moodlens")

    def gradio_predict(text):
        result = pipeline.predict(text)
        return f"Prediction: {result['label']}\n"

    iface = gr.Interface(
        fn=gradio_predict,
        inputs=gr.Textbox(lines=7, placeholder="Enter text here..."),
        outputs="text",
        title="MoodLens – Service-Focused Sentiment Analysis Agent",
        description="Moodlens is designed to evaluate the quality of service provided by the council...",
        examples=[
            "The new Customer service portal is fantastic! The new images, categories and website data model is outstanding.",
            "I had a horrible experience with the new council website upgrade. It keeps crashing and the customer support is unhelpful.",
            "The public library had great ambiance but the food was mediocre at best.",
            "I'm extremely dissatisfied with the coumcil services, contact center rep wasn't very friendly and wasn't enough knowledgeable. ",
            "Your genius pothole touch has morphed our road into an awful maze of misery, where each dreadful crater delivers a horrible shock to our wheels. We’re buzzing with the challenge, of course, but filling those gaps might cut the grief and keep us rolling smoothly.",
            "Thanks for the prompt response! My complaint has been open for three days, and my SLA is already past due. Really impressive turnaround time, you guys are at least faster than a crippled snail. Thumbs up!",
            "I'm extremely frustrated by the constant loud noises from my neighborhood—it's unbearable and needs urgent attention.",
            "Bravo on the swift response! Three days later and my issue still hasn't been touched it’s amazing how you manage to turn every complaint into a prolonged waiting experience."
        ],
    )

    iface.launch()