import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import json
from typing import Tuple

MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
ADAPTER_REPO = "riccardomusmeci/SentimentProfAI"

SYSTEM_PROMPT = """<|system|>
Analyze the sentiment of the following movie review and label it as positive or negative.
Provide ONLY an output in JSON format with two fields:
- "label": "positive" or "negative"
- "reasoning": a brief explanation of your classification

Do not add any other text after the JSON.</s>"""

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

base_model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16 if device in ["cuda", "mps"] else torch.float32,
)
tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO)
model = PeftModel.from_pretrained(base_model, ADAPTER_REPO)
model.to(device)
model.eval()

def sentiment_analysis(review_text: str) -> Tuple[str, str]:
    """Analyze the sentiment of a movie review and return the label and reasoning.
    
    Args:
        review_text (str): The movie review text to analyze.
        
    Returns:
        Tuple[str, str]: A tuple containing the sentiment label ("positive" or "negative") and the reasoning for the classification.
    
    """
    prompt = f"{SYSTEM_PROMPT}<|user|>\n{review_text}</s>\n<|assistant|>\n"
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id
        )
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    # Estrai solo la parte JSON dalla risposta
    try:
        start = response.index('{')
        end = response.rindex('}') + 1
        json_str = response[start:end]
        sentiment_json = json.loads(json_str)
        label = sentiment_json.get("label", "")
        reasoning = sentiment_json.get("reasoning", "")
    except Exception:
        label = "Errore"
        reasoning = f"Impossibile estrarre il JSON. Output grezzo: {response}"
    return label, reasoning

iface = gr.Interface(
    fn=sentiment_analysis,
    inputs=gr.Textbox(label="Movie Review"),
    outputs=[
        gr.Textbox(label="Label"),
        gr.Textbox(label="Reasoning")
    ],
    title="Sentiment Analysis ProfAI",
    description="Analizza il testo e restituisce il sentiment (positivo/negativo) e la motivazione."
)

if __name__ == "__main__":
    iface.launch()