from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer
import onnxruntime as ort
import numpy as np
from pathlib import Path
import traceback

# === Inisialisasi FastAPI ===
app = FastAPI(title="Portfolio Chatbot API", version="1.0")

# === CORS (boleh dibatasi ke domain Vercel kamu nanti) ===
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # contoh: ["https://your-frontend.vercel.app"]
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# === Path model dan tokenizer ===
BASE_DIR = Path(__file__).resolve().parent
MODEL_PATH = BASE_DIR / "models" / "bert_chatbot.onnx"
TOKENIZER_PATH = BASE_DIR / "models" / "bert-base-multilingual-cased"

# === Global variable untuk model dan tokenizer ===
tokenizer = None
session = None

# === Load model dan tokenizer ===
def load_model():
    global tokenizer, session
    try:
        print("🚀 Loading tokenizer dan ONNX model...")
        tokenizer = AutoTokenizer.from_pretrained(str(TOKENIZER_PATH))
        session = ort.InferenceSession(str(MODEL_PATH), providers=["CPUExecutionProvider"])
        print("✅ Model dan tokenizer berhasil dimuat!")
        print("📥 Model expects inputs:", [i.name for i in session.get_inputs()])
    except Exception as e:
        print("❌ ERROR saat memuat model/tokenizer:", e)
        traceback.print_exc()

load_model()

# === Label mapping (HARUS sama seperti saat training) ===
id2label = {
    0: "about_me",
    1: "career_goal",
    2: "experience",
    3: "fallback",
    4: "greeting",
    5: "projects",
    6: "skills",
}

# === Kamus respon sesuai training ===
responses = {
    "about_me": "I am a passionate developer specializing in AI and web development.",
    "skills": "My main skills are HTML5, CSS3, JavaScript, Laravel, Node.js, Database, TensorFlow, PyTorch, Firebase, and Jupyter Notebook.",
    "projects": "Some of my projects are Mobile Apps Bald Detection and Jupyter Notebook Bald Detection.",
    "experience": "I have worked as IT Support, AI Engineer, and Freelancer on multiple projects.",
    "career_goal": "My career goal is to become a Full Stack Developer and Machine Learning Engineer.",
    "greeting": "Hello! How can I help you regarding this portfolio?",
    "fallback": "I'm sorry, I don't understand. Please ask another question."
}

# === Request schema ===
class ChatRequest(BaseModel):
    text: str

# === Root endpoint ===
@app.get("/")
async def root():
    return {"message": "🚀 Portfolio Chatbot API is running successfully!"}

# === Chatbot endpoint ===
@app.post("/chatbot")
async def chatbot(req: ChatRequest):
    """
    Endpoint utama untuk memproses input teks dan mengembalikan intent serta respon.
    """
    intent = "fallback"

    # Pastikan model sudah termuat
    if session is None or tokenizer is None:
        return {"reply": responses["fallback"], "intent": "error_loading"}

    try:
        # === Tokenisasi input ===
        inputs = tokenizer(
            req.text,
            return_tensors="np",  # output dalam format numpy
            padding=True,
            truncation=True,
            max_length=128
        )

        # === Siapkan input sesuai nama yang diminta oleh model ===
        expected_inputs = [i.name for i in session.get_inputs()]
        ort_inputs = {k: v.astype(np.int64) for k, v in inputs.items() if k in expected_inputs}

        # === Jalankan inferensi ONNX ===
        ort_outputs = session.run(None, ort_inputs)
        logits = ort_outputs[0]

        # === Prediksi intent ===
        pred_id = int(np.argmax(logits, axis=1)[0])
        intent = id2label.get(pred_id, "fallback")

        # === Ambil respon ===
        reply = responses.get(intent, responses["fallback"])

        print(f"🧠 Input: {req.text} | Intent: {intent} | Reply: {reply}")

        return {"reply": reply, "intent": intent}

    except Exception as e:
        print("❌ Runtime error:", e)
        traceback.print_exc()
        return {"reply": "⚠️ Internal server error.", "intent": intent}