from __future__ import annotations

from dataclasses import dataclass, field
from typing import Literal, TypeAlias

from torch import Tensor

QuestionType: TypeAlias = Literal["choice", "score", "noul"]


@dataclass(frozen=True)
class Criterion:
    id: str
    description: str


@dataclass(frozen=True)
class Question:
    id: str
    type: QuestionType
    instructions: str
    criteria: list[Criterion | str | dict[str, str]] = field(default_factory=list)

    def normalized_criteria(self) -> list[Criterion]:
        result: list[Criterion] = []
        for index, item in enumerate(self.criteria):
            if isinstance(item, Criterion):
                result.append(item)
            elif isinstance(item, str):
                result.append(Criterion(str(index), item))
            else:
                result.append(Criterion(item.get("id", str(index)), item["description"]))
        if self.type in ("choice", "score") and not result:
            raise ValueError(f"Question {self.id!r} of type {self.type!r} needs criteria")
        return result


@dataclass
class TaskOutput:
    question_id: str
    type: QuestionType
    logits: object
    probabilities: Tensor
    prediction: object
    expected_score: object | None = None