import gradio as gr
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer

# Load model AnggaPuspa/aiboss
MODEL_NAME = "AnggaPuspa/aiboss"

print(f"Loading model: {MODEL_NAME}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

# Use max_length dari web_config (128)
MAX_LENGTH = 128
classifier = pipeline(
    "text-classification", 
    model=model, 
    tokenizer=tokenizer, 
    return_all_scores=True,
    truncation=True,
    max_length=MAX_LENGTH
)
print(f"Model loaded! Max length: {MAX_LENGTH}")

# Label mapping sesuai web_config.json - IndoBERT Sentiment Sawit
# 0 = negative, 1 = neutral, 2 = positive
LABEL_MAP = {
    "LABEL_0": "negative",
    "LABEL_1": "neutral", 
    "LABEL_2": "positive"
}

def analyze_sentiment(text: str):
    """Analyze sentiment of single input text"""
    if not text or not text.strip():
        return {"error": "Please provide text to analyze"}
    
    results = classifier(text)[0]
    
    # Normalize labels
    normalized = []
    for item in results:
        label = LABEL_MAP.get(item["label"], item["label"])
        normalized.append({"label": label, "score": round(item["score"], 4)})
    
    # Sort by score
    normalized.sort(key=lambda x: x["score"], reverse=True)
    top = normalized[0]
    
    return {
        "sentiment": top["label"],
        "confidence": top["score"],
        "details": normalized
    }

def analyze_batch(texts_json: str):
    """Analyze sentiment of multiple texts at once (batch processing)
    
    Args:
        texts_json: JSON string of array of texts, e.g. '["text1", "text2", ...]'
    
    Returns:
        List of results with sentiment analysis for each text
    """
    import json
    
    try:
        texts = json.loads(texts_json)
    except:
        return {"error": "Invalid JSON input. Expected array of strings."}
    
    if not isinstance(texts, list) or len(texts) == 0:
        return {"error": "Please provide an array of texts"}
    
    # Limit batch size to prevent timeout
    if len(texts) > 100:
        return {"error": f"Batch size too large ({len(texts)}). Max 100 texts per batch."}
    
    # Clean and truncate texts
    clean_texts = [str(t).strip()[:MAX_LENGTH * 4] if t else "" for t in texts]  # Approx char limit
    
    # Run batch inference
    try:
        all_preds = classifier(clean_texts)
        
        results = []
        for preds in all_preds:
            normalized = []
            for item in preds:
                label = LABEL_MAP.get(item["label"], item["label"])
                normalized.append({"label": label, "score": round(item["score"], 4)})
            
            normalized.sort(key=lambda x: x["score"], reverse=True)
            top = normalized[0]
            
            results.append({
                "sentiment": top["label"],
                "confidence": top["score"],
            })
        
        return results
        
    except Exception as e:
        return {"error": f"Batch processing failed: {str(e)}"}

# Create Gradio interface
with gr.Blocks(title="IndoBERT Sentiment Sawit") as demo:
    gr.Markdown("# 🌴 Sentiment Analysis API - IndoBERT Sentiment Sawit")
    gr.Markdown("Model: `AnggaPuspa/aiboss` | Labels: Negative / Neutral / Positive")
    
    with gr.Tab("Single Text"):
        single_input = gr.Textbox(
            label="Input Text", 
            placeholder="Masukkan teks untuk dianalisis...",
            lines=3
        )
        single_output = gr.JSON(label="Result")
        single_btn = gr.Button("🔍 Analyze", variant="primary")
        single_btn.click(fn=analyze_sentiment, inputs=single_input, outputs=single_output, api_name="predict")
        
        gr.Examples(
            examples=[
                "Sawit sangat bagus untuk ekonomi Indonesia",
                "Harga sawit terus menurun, petani rugi besar",
                "Sawit adalah komoditas yang biasa saja"
            ],
            inputs=single_input
        )
    
    with gr.Tab("Batch Processing"):
        batch_input = gr.Textbox(
            label="Texts (JSON Array)", 
            placeholder='["text1", "text2", "text3"]',
            lines=5
        )
        batch_output = gr.JSON(label="Results")
        batch_btn = gr.Button("Analyze Batch", variant="primary")
        batch_btn.click(fn=analyze_batch, inputs=batch_input, outputs=batch_output, api_name="predict_batch")

demo.launch()