from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import torch

app = FastAPI(title="Forex Sentiment API", version="1.0")

# ===============================
# Load Models
# ===============================
finbert_name = "ProsusAI/finbert"
longformer_name = "Miruzen/LongFormer_Skripsi"

device = 0 if torch.cuda.is_available() else -1

print("📥 Loading FinBERT model...")
finbert = pipeline(
    "text-classification",
    model=finbert_name,
    tokenizer=finbert_name,
    return_all_scores=True,
    device=device,
)

print("📥 Loading LongFormer model...")
longformer = pipeline(
    "text-classification",
    model=longformer_name,
    tokenizer=longformer_name,
    return_all_scores=True,
    device=device,
)


class InputData(BaseModel):
    title: str | None = None
    content: str | None = None


def extract_scores(predictions):
    """Convert HF model output into {positive, neutral, negative} dict."""
    scores = {"positive": 0.0, "neutral": 0.0, "negative": 0.0}
    for item in predictions[0]:
        label = item["label"].lower()
        if "pos" in label:
            scores["positive"] = item["score"]
        elif "neg" in label:
            scores["negative"] = item["score"]
        elif "neu" in label:
            scores["neutral"] = item["score"]
    dominant = max(scores, key=scores.get)
    return {"label": dominant, "scores": scores}


@app.post("/analyze")
def analyze(data: InputData):
    title_result, content_result = None, None

    if data.title:
        finbert_out = finbert(data.title)
        title_result = extract_scores(finbert_out)

    if data.content:
        longformer_out = longformer(data.content)
        content_result = extract_scores(longformer_out)

    mood_score = (
        (title_result["scores"].get("positive", 0) if title_result else 0)
        + (content_result["scores"].get("positive", 0) if content_result else 0)
        - (title_result["scores"].get("negative", 0) if title_result else 0)
        - (content_result["scores"].get("negative", 0) if content_result else 0)
    )

    return {
        "status": "ok",
        "mood_score": mood_score,
        "details": {     # 🔹 <── inilah kunci penting yang ditunggu Supabase!
            "title": title_result,
            "content": content_result,
        },
    }


@app.get("/")
def root():
    return {"message": "Forex Sentiment API active!"}