"""Unified sample schema for EvalHarness.

The data layer only *declares*; execution layers (sandbox / model / scorer)
consume these models. Raw dataset formats are free-form -- every dataset
plugin converts its records into ``Sample`` via ``record_to_sample``.
"""

from typing import Any, Dict, List, Literal, Optional, Union

from pydantic import BaseModel, Field


class ChatMessage(BaseModel):
    role: Literal['system', 'user', 'assistant', 'tool']
    content: str


class SandboxSpec(BaseModel):
    """Execution environment carried by a sample (coding/agent tasks)."""

    image: Optional[str] = None
    compose_file: Optional[str] = None
    platform: Optional[str] = None
    config: Dict[str, Any] = Field(default_factory=dict)


class ToolInfo(BaseModel):
    """Tool declaration for function-calling / agent samples."""

    name: str
    description: Optional[str] = None
    parameters: Dict[str, Any] = Field(default_factory=dict)


class Sample(BaseModel):
    """The single currency of the data layer.

    Conventions:
    - ``input``   : question text, or a list of ChatMessage for multi-turn/multimodal.
    - ``choices`` : option *contents* for multiple-choice tasks.
    - ``target``  : reference answer; a LETTER (e.g. 'A') for MCQ, text otherwise.
    - ``id``/``group_id`` : assigned by the framework on materialize/repeats.
    """

    input: Union[str, List[ChatMessage]]
    choices: Optional[List[str]] = None
    target: Union[str, List[str]] = ''
    id: Optional[int] = None
    group_id: Optional[int] = None
    task_type: Optional[str] = None  # qa | mcq | math | coding | agent | vqa | ...
    tools: Optional[List[ToolInfo]] = None
    sandbox: Optional[SandboxSpec] = None
    files: Optional[Dict[str, str]] = None  # path -> content, copied into sandbox
    setup: Optional[str] = None  # script run in sandbox before use
    metadata: Dict[str, Any] = Field(default_factory=dict)

    @property
    def input_text(self) -> str:
        """Unified text view of the input."""
        if isinstance(self.input, str):
            return self.input
        return '\n'.join(m.content for m in self.input)