import diceplayer.dice.dice_input as dice_input
from diceplayer.config import DiceConfig

from pydantic import TypeAdapter
from pydantic.dataclasses import dataclass

import subprocess
from itertools import islice
from pathlib import Path
from typing import Final, List, Self


@dataclass(slots=True, frozen=True)
class DiceEnvironmentItem:
    atom: str
    x: float
    y: float
    z: float


DiceEnvironmentItemAdapter = TypeAdapter(DiceEnvironmentItem)


@dataclass(slots=True)
class DiceEnvironment:
    number_of_sites: int
    thickness: List[float]
    items: List[DiceEnvironmentItem]

    @classmethod
    def new(cls, thickness: List[float]) -> Self:
        return cls(number_of_sites=0, thickness=thickness, items=[])

    def add_site(self, site: DiceEnvironmentItem):
        self.items.append(site)
        self.number_of_sites += 1


DICE_FLAG_LINE: Final[int] = -2
DICE_END_FLAG: Final[str] = "End of simulation"


class DiceWrapper:
    def __init__(self, dice_config: DiceConfig, working_directory: Path):
        self.dice_config = dice_config
        self.working_directory = working_directory

    def run(self, dice_config: dice_input.DiceInputConfig) -> None:
        input_path = dice_input.write_config(dice_config, self.working_directory)
        output_path = input_path.parent / (input_path.name + ".out")

        with open(output_path, "w") as outfile, open(input_path, "r") as infile:
            exit_status = subprocess.call(
                self.dice_config.progname,
                stdin=infile,
                stdout=outfile,
                cwd=self.working_directory,
            )

        if exit_status != 0:
            raise RuntimeError(f"Dice simulation failed with exit status {exit_status}")

        with open(output_path, "r") as outfile:
            line = outfile.readlines()[DICE_FLAG_LINE]
            if line.strip() == DICE_END_FLAG:
                return

        raise RuntimeError(f"Dice simulation failed with exit status {exit_status}")

    def parse_results(self) -> list[DiceEnvironment]:
        results = []

        positions_file = self.working_directory / "phb.xyz"
        if not positions_file.exists():
            raise RuntimeError(f"Positions file not found at {self.working_directory}")

        with open(positions_file, "r") as f:
            while True:
                line = f.readline()
                if not line.startswith("    "):
                    break

                environment = DiceEnvironment(
                    number_of_sites=int(line.strip()),
                    thickness=[float(n) for n in f.readline().split()[-3:]],
                    items=[],
                )

                # Skip the comment line

                environment.items.extend(
                    [
                        DiceEnvironmentItemAdapter.validate_python(
                            {"atom": site[0], "x": site[1], "y": site[2], "z": site[3]}
                        )
                        for atom in islice(f, environment.number_of_sites)
                        if (site := atom.strip().split()) and len(site) == 4
                    ]
                )

                results.append(environment)

        return results