import torch
import emoji
import re
from transformers import BertTokenizer, BertForSequenceClassification
from fastapi import FastAPI
from pydantic import BaseModel
from typing import Dict, List

# ====================================================================
# 1. KELAS LOGIKA ANDA (Tidak ada perubahan)
# ====================================================================

# Impor library yang dibutuhkan
import re # Untuk operasi regular expression
import emoji # Untuk menangani emoji

# Definisikan sebuah kelas untuk mengelompokkan semua fungsi pembersihan teks
class TextCleaner:
    # Metode constructor, dieksekusi saat objek TextCleaner dibuat
    def __init__(self):
        # Inisialisasi daftar karakter yang akan diperiksa untuk pengulangan.
        # Daftar ini berisi berbagai tanda baca dan simbol.
        self.character = ['.', ',', ';', ':', '?', '!', '(', ')', '[', ']', '{', '}', '<', '>', '"', '/', '\'', '-', '@']
        # Tambahkan semua huruf abjad (a-z) ke dalam daftar karakter di atas.
        self.character.extend([chr(i) for i in range(ord('a'), ord('z') + 1)])

    # Metode untuk membersihkan karakter yang berulang lebih dari 2 kali (misal: "haaiiii" -> "haai").
    def repeatcharClean(self, text):
        # Ulangi untuk setiap karakter dalam daftar 'self.character'
        for char_to_clean in self.character:
            # Buat pola regex untuk menemukan karakter yang berulang 3 kali atau lebih secara berurutan.
            # Contoh: jika char_to_clean adalah 'a', polanya akan mencari 'aaa' atau 'aaaa', dst.
            pattern = re.compile(re.escape(char_to_clean) + r'{3,}')
            # Ganti urutan karakter yang berulang tersebut dengan satu karakter saja.
            # Contoh: "good moooorning" menjadi "good morning".
            text = pattern.sub(char_to_clean, text)
        return text

    # Metode utama untuk menjalankan seluruh proses pembersihan teks
    def clean_review(self, text):
        # 1. Ubah semua teks menjadi huruf kecil (lowercase) untuk konsistensi.
        text = text.lower()
        # 2. Ganti spasi, tab, atau baris baru yang berlebih dengan satu spasi saja.
        text = re.sub(r'\s+', ' ', text)
        # 3. Hapus karakter non-ASCII (seperti karakter Cina, Arab, atau beberapa emoji kompleks).
        text = re.sub(r'[^\x00-\x7F]+', ' ', text)

        # 4. Ganti @mention dan URL dengan token placeholder.
        new_text = []
        # Pecah teks menjadi kata-kata
        for word in text.split(" "):
            # Jika kata diawali dengan '@' dan lebih dari 1 karakter, ganti dengan '@USER'.
            word = '@USER' if word.startswith('@') and len(word) > 1 else word
            # Jika kata diawali dengan 'http', ganti dengan 'HTTPURL'.
            word = 'HTTPURL' if word.startswith('http') else word
            new_text.append(word)
        # Gabungkan kembali kata-kata menjadi satu kalimat.
        text = " ".join(new_text)

        # 5. Ubah emoji menjadi representasi teksnya (misal: 😊 -> ':smiling_face:').
        text = emoji.demojize(text)
        # 6. Hapus representasi teks emoji yang polanya seperti ':nama_emoji:'.
        text = re.sub(r':[A-Za-z_-]+:', ' ', text)

        # 7. Hapus emoticon umum berbasis teks seperti :), :D, :(, xD, dll.
        text = re.sub(r"([xX;:]'?[dDpPvVoO3)(])", ' ', text)
        # 8. Hapus semua tanda baca dan simbol yang tersisa.
        text = re.sub(r'["#$%&()*+,./:;<=>\[\]\\^_`{|}~]', ' ', text)
        
        # 9. Panggil metode 'repeatcharClean' untuk membersihkan karakter yang berulang.
        text = self.repeatcharClean(text)
        
        # 10. Lakukan pembersihan spasi terakhir dan hapus spasi di awal/akhir kalimat.
        text = re.sub(r'\s+', ' ', text).strip()
        
        # Kembalikan teks yang sudah bersih
        return text

class SentimentPredictor:
    def __init__(self, tokenizer, model):
        self.tokenizer = tokenizer
        self.model = model
        self.device = torch.device("cpu")
        self.model.to(self.device)
        self.label_mapping = {0: 'Positif', 1: 'Netral', 2: 'Negatif'}

    def predict(self, text: str) -> (str, float, Dict[str, float]):
        inputs = self.tokenizer(text, return_tensors='pt', truncation=True, padding=True, max_length=280)
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = self.model(**inputs)
        
        logits = outputs.logits
        probabilities = torch.softmax(logits, dim=1)[0]
        
        confidence_score = probabilities.max().item()
        predicted_label_id = probabilities.argmax().item()
        sentiment = self.label_mapping[predicted_label_id]
        
        all_scores = {self.label_mapping[i]: prob.item() for i, prob in enumerate(probabilities)}
        return sentiment, confidence_score, all_scores

# ====================================================================
# 2. INISIALISASI MODEL & APLIKASI FASTAPI (Tidak ada perubahan)
# ====================================================================

print("Memuat model dan tokenizer...")
tokenizer = BertTokenizer.from_pretrained('indolem/indobertweet-base-uncased')
model = BertForSequenceClassification.from_pretrained('indolem/indobertweet-base-uncased', num_labels=3)
model_path = 'model_indoBERTweet_100Epochs_sentiment.pth'
state_dict = torch.load(model_path, map_location=torch.device('cpu'))
model.load_state_dict(state_dict, strict=False)
model.eval()
print("Model berhasil dimuat.")

text_cleaner = TextCleaner()
sentiment_predictor = SentimentPredictor(tokenizer, model)

app = FastAPI(
    title="API Klasifikasi Sentimen",
    description="Sebuah API untuk menganalisis sentimen teks Bahasa Indonesia."
)

# ====================================================================
# 3. DEFINISIKAN MODEL INPUT & OUTPUT API
# ====================================================================

class TextInput(BaseModel):
    text: str

# --- [PERBAIKAN] --- Menambahkan definisi BatchTextInput ---
# Model ini memberitahu FastAPI bahwa endpoint batch akan menerima
# sebuah objek JSON dengan satu key "texts" yang berisi daftar string.
class BatchTextInput(BaseModel):
    texts: List[str]
# -----------------------------------------------------------

class PredictionOutput(BaseModel):
    cleaned_text: str = None  # Optional, hanya diisi pada batch
    sentiment: str
    confidence: float
    all_scores: Dict[str, float]

# ====================================================================
# 4. BUAT ENDPOINT (Tidak ada perubahan logika)
# ====================================================================

@app.get("/")
def read_root():
    return {"message": "Selamat datang di API Klasifikasi Sentimen"}

@app.post("/predict", response_model=PredictionOutput)
def predict_sentiment(request: TextInput):
    cleaned_text = text_cleaner.clean_review(request.text)
    sentiment, confidence, all_scores = sentiment_predictor.predict(cleaned_text)
    return PredictionOutput(
        sentiment=sentiment, 
        confidence=confidence, 
        all_scores=all_scores
    )

@app.post("/predict-batch", response_model=List[PredictionOutput])
def predict_sentiment_batch(request: BatchTextInput):
    results = []
    for text in request.texts:
        cleaned_text = text_cleaner.clean_review(text)
        sentiment, confidence, all_scores = sentiment_predictor.predict(cleaned_text)
        results.append(PredictionOutput(
            cleaned_text=cleaned_text,
            sentiment=sentiment, 
            confidence=confidence, 
            all_scores=all_scores
        ))
    return results