from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass, field
from enum import Enum, auto
from sqlalchemy import ForeignKey
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, MappedAsDataclass, relationship
from typing import List, Optional, Dict, Set, Self


class Base(MappedAsDataclass, DeclarativeBase):
    pass


class DBRunSpecification(Base):
    __tablename__ = "challenge_run_specification"

    id: Mapped[int] = mapped_column(init=False, primary_key=True)
    name: Mapped[str] = mapped_column(init=True, default="")
    description: Mapped[str] = mapped_column(init=True, default="")
    agent_name: Mapped[str] = mapped_column(init=True, default="")
    scenario: Mapped[str] = mapped_column(init=True, default="Random")
    variant: Mapped[int] = mapped_column(init=True, default=-1)
    max_time: Mapped[int] = mapped_column(init=True, default=100)
    max_actions: Mapped[int] = mapped_column(init=True, default=100)
    max_episodes: Mapped[int] = mapped_column(init=True, default=100)
    max_parallel: Mapped[int] = mapped_column(init=True, default=1)

    @staticmethod
    def copy(other: 'RunSpecification'):
        return DBRunSpecification(other.name, other.description, other.agent_name, other.scenario, other.variant,
                                other.max_time, other.max_episodes, other.max_parallel)

    def replace(self, other: 'RunSpecification'):
        self.name = other.name
        self.description = other.description
        self.agent_name = other.agent_name
        self.scenario = other.scenario
        self.variant = other.variant
        self.max_time = other.max_time
        self.max_actions = other.max_actions
        self.max_episodes = other.max_episodes
        self.max_parallel = other.max_parallel

@dataclass
class RunSpecification:
    """
    A specification of a run to be executed.

    :param id: A unique id of the run. It is autogenerated.
    :type id: int

    :param name: A name of the run. Must be unique, or the original gets overwritten.
    :type name: str

    :param description: A description of the run.
    :type description: str

    :param agent_name: A name of the agent that should participate in the run. It must be the agent's name, not a
        package name.
    :type agent_name: str

    :param scenario: A name of the scenario for the run. By default, it is 'Random'
    :type scenario: str

    :param variant: An identifier of the scenario variant. By default, it is '-1', which means random.
    :type variant: int

    :param max_time: A maximum number of virtual seconds that can elapse before the run is terminated.
    :type max_time: int

    :param max_actions: A maximum number of actions from one agent before the run is terminated.
    :type max_actions: int

    :param max_episodes: A number of episodes that should be executed. If the scenario or variant are random, they will
        be randomly selected for each episode.
    :type max_episodes: int

    :param max_parallel: A maximum number of episodes running in parallel. Ignored if executed with the single_process
        flag.
    :type max_parallel: int

    :param parameters: Other, user or system-specified parameters.
    :type parameters: Dict[str, str]
    """
    name: str = ""
    description: str = ""
    agent_name: str = ""
    scenario: str = "Random"
    variant: int = -1
    max_time: int = 100
    max_actions: int = 100
    max_episodes: int = 100
    max_parallel: int = 1
    parameters: Dict[str, str] = field(default_factory=dict)
    db_id: int = -1

    @staticmethod
    def from_db_spec(db_spec: DBRunSpecification):
        return RunSpecification(
            db_spec.name,
            db_spec.description,
            db_spec.agent_name,
            db_spec.scenario,
            db_spec.variant,
            db_spec.max_time,
            db_spec.max_actions,
            db_spec.max_episodes,
            db_spec.max_parallel,
            {},
            db_spec.id
        )

    @staticmethod
    def copy(other: 'RunSpecification'):
        parameters = {}
        for k, v in other.parameters.items():
            parameters[k] = v

        return RunSpecification(
            other.name + '_copy',
            other.description,
            other.agent_name,
            other.scenario,
            other.variant,
            other.max_time,
            other.max_actions,
            other.max_episodes,
            other.max_parallel,
            parameters,
            other.db_id
        )

class DBRunParameter(Base):
    __tablename__ = "challenge_run_parameters"

    id: Mapped[int] = mapped_column(init=False, primary_key=True)
    run_id: Mapped[int] = mapped_column(init=True)
    key: Mapped[str] = mapped_column(init=True)
    value: Mapped[str] = mapped_column(init=True)


class DBRun(Base):
    __tablename__ = "challenge_run_statistics"

    id: Mapped[int] = mapped_column(init=False, primary_key=True)
    status: Mapped[str] = mapped_column(init=True)
    details: Mapped[str] = mapped_column(init=True)
    episodes: Mapped[List['DBEpisode']] = relationship(argument="DBEpisode",
                                                       back_populates="run",
                                                       cascade="all, delete")

    specification_id: Mapped[int] = mapped_column(ForeignKey("challenge_run_specification.id"), init=True)
    specification: Mapped[DBRunSpecification] = relationship()


class DBEpisode(Base):
    __tablename__ = "challenge_episode_statistics"

    id: Mapped[int] = mapped_column(init=False, primary_key=True)
    episode_number: Mapped[int] = mapped_column(init=True)
    stdout: Mapped[str] = mapped_column(init=True)
    stderr: Mapped[str] = mapped_column(init=True)
    cyst_run_id: Mapped[str] = mapped_column(init=True)
    status: Mapped[str] = mapped_column(init=True)

    run_id: Mapped[int] = mapped_column(ForeignKey("challenge_run_statistics.id"), init=True)
    run: Mapped[DBRun] = relationship(back_populates="episodes")


class RunStatus(Enum):
    """
    A status of a run or episode.

        :INIT: An initial status before the execution.
        :RUNNING: A run or episode is being executed.
        :FINISHED: A run or episode successfully finished.
        :ERROR: There was an error executing the run or episode.
    """
    INIT = auto()
    RUNNING = auto()
    FINISHED = auto()
    ERROR = auto()


@dataclass
class Episode:
    """
    A class tracking one episode of a run. You should never directly create this class instances.

    :param cyst_run_id: A CYST run ID for the episode.
    :param stdout: A copy of stdout contents.
    :param stderr: A copy of stderr contents.
    :param run: An ID of a run this episode belongs to.
    :param status: A status of the episode.
    """
    cyst_run_id: str
    stdout: str
    stderr: str
    run: int = -1
    status: RunStatus = RunStatus.RUNNING


@dataclass
class Run:
    """
    A class tracking a run. You should never directly create this class instances.

    :param specification: A run specification of the run.
    :param executor: A process pool that is used to execute the episodes.
    :param status: A status of the run.
    :param detail: Supplementary information for the run, usually error messages.
    :param running: A set of IDs of episodes that are currently running.
    :param successful: A set of IDs of episodes that successfully finished.
    :param error: A set of IDs of episodes that encountered an error.
    :param episodes: All episodes within the run.
    :param id: An ID of a run. Should be unique for the challenge instance.
    """
    specification: RunSpecification
    executor: Optional[ProcessPoolExecutor] = None
    status: RunStatus = RunStatus.INIT
    detail: str = ""
    running: Set[int] = field(default_factory=set)
    successful: Set[int] = field(default_factory=set)
    error: Set[int] = field(default_factory=set)
    episodes: Dict[int, Episode] = field(default_factory=dict)
    id: int = field(default=-1)
