from typing import Any, Optional

from pydantic import model_validator

from galtea.utils.from_camel_case_base_model import FromCamelCaseBaseModel
from galtea.utils.input_normalization import coerce_to_string, unwrap_context, unwrap_input


class TestCaseBase(FromCamelCaseBaseModel):
    """
    Base model for a test case.

    Attributes:
        test_id (str): ID of the test.
        input (Optional[str]): Input for the test case.
        expected_output (Optional[str]): Expected output for the test case.
        context (Optional[str]): Context for the test case.
        expected_tools (Optional[list[str]]): Expected tools for the test case evaluation.
        source (Optional[str]): Source of the test case.
        strategy (Optional[str]): Strategy for the test case.
        variant (Optional[str]): Variant for the test case.
        reviewed_by_id (Optional[str]): ID of the user who reviewed the test case.
        language_code (Optional[str]): Language code for the test case.
        user_score (Optional[int]): User score for the test case.
        user_score_reason (Optional[str]): User score reason for the test case.
        input_data (Optional[Dict]): Raw JSON input object from the API. Use ``input`` for
            the backward-compatible string value.
        context_data (Optional[Dict]): Raw JSON context object from the API. Use ``context``
            for the backward-compatible string value.
    """

    test_id: str

    # ACCURACY / SAFETY Test Case
    input: Optional[str] = None
    expected_output: Optional[str] = None
    context: Optional[str] = None
    expected_tools: Optional[list[str]] = None
    variant: Optional[str] = None
    strategy: Optional[str] = None

    # BEHAVIOR Test Case
    user_persona: Optional[str] = None
    scenario: Optional[str] = None
    goal: Optional[str] = None
    max_iterations: Optional[int] = None
    stopping_criterias: Optional[list[str]] = None

    # COMMON FIELDS
    source: Optional[str] = None
    is_augmented: Optional[bool] = False
    reviewed_by_id: Optional[str] = None
    language_code: Optional[str] = None
    user_score: Optional[int] = None
    user_score_reason: Optional[str] = None
    confidence: Optional[float] = None
    confidence_reason: Optional[str] = None

    # Raw JSON objects from the API (for structured input/context access)
    input_data: Optional[dict[str, Any]] = None
    context_data: Optional[dict[str, Any]] = None

    @model_validator(mode="before")
    @classmethod
    def _unwrap_json_fields(cls, data: Any) -> Any:
        """Auto-unwrap JSON input/context from the API into backward-compatible strings."""
        if not isinstance(data, dict):
            return data

        if "input" in data:
            raw_input = data["input"]
        elif "initialPrompt" in data:
            raw_input = data["initialPrompt"]
        else:
            raw_input = data.get("initial_prompt")
        if isinstance(raw_input, dict):
            data["input_data"] = raw_input
            data["input"] = coerce_to_string(unwrap_input(raw_input))
        elif isinstance(raw_input, str):
            data["input"] = raw_input

        raw_context = data.get("context")
        if isinstance(raw_context, dict):
            data["context_data"] = raw_context
            data["context"] = coerce_to_string(unwrap_context(raw_context))
        elif isinstance(raw_context, str):
            data["context"] = raw_context

        return data


class TestCase(TestCaseBase):
    """
    Model for a test case, including database identifiers and timestamps.

    Attributes:
        id (str): Unique identifier for the test case.
        created_at (str): Creation timestamp.
        deleted_at (Optional[str]): Deletion timestamp, if deleted.
    """

    id: str
    created_at: str
    deleted_at: Optional[str] = None
