from typing import TYPE_CHECKING, Any, Literal

import numpy as np
from pydantic import Field, model_validator

from ..base import BaseModel

if TYPE_CHECKING:
    import torch


class TensorData(BaseModel):
    data: list[float] | list[int]

    dtype: Literal["float32", "int64"]

    shape: list[int] = Field(default_factory=list)

    @model_validator(mode="before")
    @classmethod
    def convert_input(cls, value: Any) -> Any:
        if isinstance(value, cls):
            return value

        try:
            import torch

            if isinstance(value, torch.Tensor):
                return cls.from_torch(value).model_dump()
        except ImportError:
            pass

        if isinstance(value, np.ndarray):
            return cls.from_numpy(value).model_dump()

        if isinstance(value, list):
            arr = np.asarray(value)

            if arr.dtype.kind not in ("i", "u", "f"):
                raise ValueError("Only numeric lists are allowed for TensorData")

            dtype: Literal["float32", "int64"]
            if arr.dtype.kind == "f":
                arr = arr.astype(np.float32)
                dtype = "float32"
            else:
                arr = arr.astype(np.int64)
                dtype = "int64"

            return {
                "data": arr.reshape(-1).tolist(),
                "dtype": dtype,
                "shape": list(arr.shape),
            }

        return value

    @model_validator(mode="after")
    def validate_tensor_data(self) -> "TensorData":
        if not self.shape:
            self.shape = [len(self.data)]

        if any(dim < 0 for dim in self.shape):
            raise ValueError("Negative dimensions are not allowed in shape")

        expected = 1
        for dim in self.shape:
            expected *= dim

        if expected != len(self.data):
            raise ValueError(
                f"Unexpected data length {len(self.data)} for shape {self.shape} (expected {expected})"
            )

        if self.dtype == "float32":
            self.data = [float(x) for x in self.data]
        else:
            normalized: list[int] = []
            for x in self.data:
                if isinstance(x, float) and not x.is_integer():
                    raise ValueError(
                        f"Expected integer values for dtype 'int64', got float {x}"
                    )
                normalized.append(int(x))
            self.data = normalized

        return self

    @classmethod
    def from_numpy(cls, array: np.ndarray) -> "TensorData":
        if np.issubdtype(array.dtype, np.floating):
            array = array.astype(np.float32)
            dtype: Literal["float32", "int64"] = "float32"
        elif np.issubdtype(array.dtype, np.integer):
            array = array.astype(np.int64)
            dtype = "int64"
        else:
            raise ValueError(f"Unsupported numpy dtype for TensorData: {array.dtype}")

        return cls(
            data=array.reshape(-1).tolist(),
            dtype=dtype,
            shape=list(array.shape),
        )

    def to_numpy(self) -> np.ndarray:
        if self.dtype == "float32":
            numpy_dtype = np.float32
        elif self.dtype == "int64":
            numpy_dtype = np.int64
        else:
            raise ValueError(f"Unsupported TensorDtype: {self.dtype}")

        arr = np.array(self.data, dtype=numpy_dtype)
        if self.shape:
            arr = arr.reshape(self.shape)
        return arr

    @classmethod
    def from_torch(cls, tensor: "torch.Tensor") -> "TensorData":
        try:
            import torch
        except ImportError as e:
            raise ImportError(
                "Detected missing 'torch'. This operation requires it. "
                f"Install via 'pip install torch'. Error details: {e}"
            )

        if tensor.dtype.is_floating_point:
            tensor = tensor.to(torch.float32)
            dtype: Literal["float32", "int64"] = "float32"
        else:
            tensor = tensor.to(torch.int64)
            dtype = "int64"

        return cls(
            data=tensor.detach().cpu().reshape(-1).tolist(),
            dtype=dtype,
            shape=list(tensor.shape),
        )

    def tolist(self) -> list[Any]:
        return self.to_numpy().tolist()

    def __len__(self) -> int:
        return len(self.data)
