from transformers import pipeline
from fastapi import FastAPI, Request
import uvicorn

from uagents import Agent, Context, Bureau, Model

# ─── uAgent I/O MODEL ─────────────────────────────────────────────────────

class TextInput(Model):
    text: str

# ─── LOAD EMOTION PIPELINE ─────────────────────────────────────────────────

emotion_model = pipeline(
    "text-classification",
    model="bhadresh-savani/distilbert-base-uncased-emotion"
)

# ─── CUSTOM ANALYSIS LOGIC ─────────────────────────────────────────────────

def analyze_text_metrics(text):
    results = emotion_model(text)
    t = text.lower()

    suicide_keywords   = ["kill myself", "suicidal", "die", "ending it", "pills", "overdose", "no way out"]
    psychosis_keywords = ["voices", "hallucinate", "not real", "they’re watching me", "i’m not me"]

    metrics = {
        "self_harm": 0.0,
        "homicidal": 0.0,
        "distress":  0.0,
        "psychosis": 0.0
    }

    # aggregate raw scores
    for res in results:
        label = res["label"]
        score = res["score"]
        if label == "sadness":
            metrics["self_harm"] += score
            metrics["distress"]  += score * 0.6
        elif label in ("anger", "fear"):
            metrics["homicidal"] += score
            metrics["distress"]  += score * 0.5
        elif label == "joy":
            metrics["psychosis"] += score * 0.3
        elif label == "surprise":
            metrics["psychosis"] += score * 0.5

    # keyword overrides
    if any(kw in t for kw in suicide_keywords):
        metrics["self_harm"] = max(metrics["self_harm"], 0.8)
    if any(kw in t for kw in psychosis_keywords):
        metrics["psychosis"] = max(metrics["psychosis"], 0.8)

    # clamp into [0.01, 0.99]
    for k in metrics:
        val = metrics[k]
        clamped = max(min(val, 0.99), 0.01)
        metrics[k] = round(clamped, 2)

    return metrics

# ─── uAgent DEFINITION ────────────────────────────────────────────────────

agent = Agent(name="sentiment_agent")

@agent.on_message(model=TextInput)
async def handle_message(ctx: Context, sender: str, msg: TextInput):
    flags = analyze_text_metrics(msg.text)
    await ctx.send(sender, flags)

# ─── FASTAPI HTTP ENDPOINT ───────────────────────────────────────────────

app = FastAPI()

@app.post("/")
async def analyze_text(request: Request):
    data = await request.json()
    return analyze_text_metrics(data.get("text", ""))

# ─── START BOTH ───────────────────────────────────────────────────────────

if __name__ == "__main__":
    bureau = Bureau()
    bureau.add(agent)
    bureau.run_in_thread()            # serve the agent on Agentverse
    uvicorn.run(app, host="0.0.0.0", port=8000)