"""Custom HF Inference Endpoint handler for PIXMELT/laya (CPU, low-memory).

Replicates test.py from the model snapshot: same state + same payload,
timed server-side. Defaults kick in when the request body is empty.
"""
import time

DEFAULT_STATE = {
    "body": (
        "By 2031, artificial intelligences no longer merely answer: they anticipate, "
        "negotiate and sometimes push back. In newsrooms, hospitals and courtrooms, "
        "autonomous agents propose decisions, draft summaries and flag their own "
        "uncertainty. Some observers speak of a 'second cognitive wave': systems capable "
        "of doubting, asking for clarification and refusing a poorly specified task. "
        "Others worry about a silent dependency, where humans approve without truly "
        "understanding. One thing is certain: the question is no longer whether AIs "
        "think, but how we will coexist with them."
    )
}

DEFAULT_PAYLOAD = {
    "tone": {
        "type": "choice",
        "instructions": "What is the dominant tone of `body`?",
        "criteria": {
            "neutral": "factual account, no judgment",
            "alarmist": "emphasis on risks and dangers",
            "enthusiastic": "emphasis on promises and benefits",
        },
    },
    "domain_coverage": {
        "type": "score",
        "instructions": "How many distinct application domains are covered in `body`?",
        "criteria": ["none", "one", "two", "three or more"],
    },
    "mentions_doubt": {
        "type": "noul",
        "instructions": "Does the text mention AIs' ability to doubt or flag their own uncertainty?",
    },
}


class EndpointHandler:
    def __init__(self, path="."):
        import torch

        # Build the model directly in bf16 so weights stay ~0.85 GB (fits in 2 GB RAM).
        torch.set_default_dtype(torch.bfloat16)
        import laya

        agent = laya.load("PIXMELT/laya", device="cpu")
        try:
            agent.model.to(torch.bfloat16)
            # laya's forward does h[:, 0].float() before act_head -> keep that head fp32
            agent.model.act_head = agent.model.act_head.to(torch.float32)
        except Exception:
            pass
        agent.dtype = torch.bfloat16
        self.agent = agent

    def __call__(self, data):
        inputs = data.get("inputs", data) if isinstance(data, dict) else {}
        state = inputs.get("state", DEFAULT_STATE)
        payload = inputs.get("payload", DEFAULT_PAYLOAD)

        timer = time.perf_counter()
        result = self.agent.predict(state, payload)
        duration_ms = round((time.perf_counter() - timer) * 1000, 2)

        answers = {
            name: result["answers"][name][payload[name]["type"]]
            for name in payload
        }
        return {
            "answers": answers,
            "duration_ms": duration_ms,
            "usage": result.get("usage", {}),
        }