from diceplayer.config import PlayerConfig

from typing_extensions import Self

import random
from abc import ABC
from dataclasses import dataclass, fields
from typing import Any, Sequence, TextIO


@dataclass(slots=True)
class BaseConfig(ABC):
    ncores: int
    ljname: str
    outname: str
    nmol: Sequence[int]
    temp: float
    seed: int
    isave: int
    press: float = 1.0

    def write_dice_config(self, io_writer: TextIO) -> None:
        for field in fields(self):
            key = field.name
            value = getattr(self, key)

            if value is None:
                continue

            io_writer.write(f"{key} = {self._serialize_value(value)}\n")

        io_writer.write("$end\n")

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        base_fields = cls._extract_base_fields(config)
        return cls(**base_fields, **kwargs)

    @staticmethod
    def _extract_base_fields(config: PlayerConfig) -> dict[str, Any]:
        seed: int
        if config.dice.seed is not None:
            seed = config.dice.seed
        else:
            seed = random.randint(0, 2**32 - 1)

        return dict(
            ncores=config.ncores,
            ljname=config.dice.ljname,
            outname=config.dice.outname,
            nmol=config.dice.nmol,
            temp=config.dice.temp,
            seed=seed,
            isave=config.dice.isave,
            press=config.dice.press,
        )

    @staticmethod
    def _get_nstep(config: PlayerConfig, idx: int) -> int:
        if len(config.dice.nstep) > idx:
            return config.dice.nstep[idx]
        return config.dice.nstep[-1]

    @staticmethod
    def _serialize_value(value: Any) -> str:
        if value is None:
            raise ValueError("DICE configuration cannot serialize None values")

        if isinstance(value, bool):
            return "yes" if value else "no"

        if isinstance(value, (list, tuple)):
            return " ".join(str(v) for v in value)

        return str(value)


# -----------------------------------------------------
# NVT BASE
# -----------------------------------------------------


@dataclass(slots=True)
class NVTConfig(BaseConfig):
    title: str = "Diceplayer Run - NVT"
    dens: float = 0.0
    nstep: int = 0
    vstep: int = 0


# -----------------------------------------------------
# NVT THERMALIZATION
# -----------------------------------------------------


@dataclass(slots=True)
class NVTTerConfig(NVTConfig):
    title: str = "Diceplayer Run - NVT Thermalization"
    upbuf: int = 360

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        return super(NVTTerConfig, cls).from_config(
            config,
            dens=config.dice.dens,
            nstep=cls._get_nstep(config, 0),
            upbuf=config.dice.upbuf,
            vstep=0,
            **kwargs,
        )


# -----------------------------------------------------
# NVT PRODUCTION
# -----------------------------------------------------


@dataclass(slots=True)
class NVTEqConfig(NVTConfig):
    title: str = "Diceplayer Run - NVT Production"
    irdf: int = 0

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        return super(NVTEqConfig, cls).from_config(
            config,
            dens=config.dice.dens,
            nstep=cls._get_nstep(config, 1),
            irdf=config.dice.irdf,
            vstep=0,
            **kwargs,
        )


# -----------------------------------------------------
# NPT BASE
# -----------------------------------------------------


@dataclass(slots=True)
class NPTConfig(BaseConfig):
    title: str = "Diceplayer Run - NPT"
    nstep: int = 0
    vstep: int = 5000


# -----------------------------------------------------
# NPT THERMALIZATION
# -----------------------------------------------------


@dataclass(slots=True)
class NPTTerConfig(NPTConfig):
    title: str = "Diceplayer Run - NPT Thermalization"
    dens: float | None = None

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        return super(NPTTerConfig, cls).from_config(
            config,
            dens=config.dice.dens,
            nstep=cls._get_nstep(config, 1),
            vstep=config.dice.vstep,
            **kwargs,
        )


# -----------------------------------------------------
# NPT PRODUCTION
# -----------------------------------------------------


@dataclass(slots=True)
class NPTEqConfig(NPTConfig):
    title: str = "Diceplayer Run - NPT Production"
    dens: float | None = None

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        return super(NPTEqConfig, cls).from_config(
            config,
            dens=config.dice.dens,
            nstep=cls._get_nstep(config, 2),
            vstep=config.dice.vstep,
            **kwargs,
        )