from collections import OrderedDict
from decimal import Decimal
import enum
from typing import NamedTuple, Optional

import graphene
from graphene.types.datetime import DateTime as GQLDateTime
import sqlalchemy as sa

from .base import metadata, EnumType

__all__ = (
    'agents', 'AgentStatus', 'ResourceSlot',
    'Agent',
)


class ResourceSlot(NamedTuple):
    id: Optional[str] = None
    # following slots ( = shares) may be fractional.
    mem: Decimal = Decimal(0)  # multiple of GiBytes
    cpu: Decimal = Decimal(0)  # multiple of CPU cores
    # TODO: dynamic slots
    cuda: Decimal = Decimal(0)  # multiple of CUDA devices
    tpu: Decimal = Decimal(0)  # multiple of TPU devices


class AgentStatus(enum.Enum):
    ALIVE = 0
    LOST = 1
    RESTARTING = 2
    TERMINATED = 3


agents = sa.Table(
    # TODO: change accelerator slots to be dynamically defined.

    'agents', metadata,
    sa.Column('id', sa.String(length=64), primary_key=True),
    sa.Column('status', EnumType(AgentStatus), nullable=False, index=True,
              default=AgentStatus.ALIVE),
    sa.Column('region', sa.String(length=64), index=True, nullable=False),

    sa.Column('mem_slots', sa.BigInteger(), nullable=False),  # MiBytes
    sa.Column('cpu_slots', sa.Float(), nullable=False),  # number of CPU cores
    sa.Column('gpu_slots', sa.Float(), nullable=False),  # number of CUDA devices
    sa.Column('tpu_slots', sa.Float(), nullable=False),  # number of TPU devices

    sa.Column('used_mem_slots', sa.BigInteger(), nullable=False),
    sa.Column('used_cpu_slots', sa.Float(), nullable=False),
    sa.Column('used_gpu_slots', sa.Float(), nullable=False),
    sa.Column('used_tpu_slots', sa.Float(), nullable=False),

    sa.Column('addr', sa.String(length=128), nullable=False),
    sa.Column('first_contact', sa.DateTime(timezone=True),
              server_default=sa.func.now()),
    sa.Column('lost_at', sa.DateTime(timezone=True), nullable=True),
)


class Agent(graphene.ObjectType):
    # TODO: change accelerator slots to be dynamically defined.

    id = graphene.String()
    status = graphene.String()
    region = graphene.String()
    mem_slots = graphene.Int()
    cpu_slots = graphene.Float()
    gpu_slots = graphene.Float()
    tpu_slots = graphene.Float()
    used_mem_slots = graphene.Int()
    used_cpu_slots = graphene.Float()
    used_gpu_slots = graphene.Float()
    used_tpu_slots = graphene.Float()
    addr = graphene.String()
    first_contact = GQLDateTime()
    lost_at = GQLDateTime()

    # Dynamic fields
    cpu_cur_pct = graphene.Float()
    mem_cur_bytes = graphene.Float()

    computations = graphene.List(
        'ai.backend.manager.models.Computation',
        status=graphene.String())

    @classmethod
    def from_row(cls, row):
        if row is None:
            return None
        return cls(
            id=row['id'],
            status=row['status'],
            region=row['region'],
            mem_slots=row['mem_slots'],
            cpu_slots=row['cpu_slots'],
            gpu_slots=row['gpu_slots'],
            tpu_slots=row['tpu_slots'],
            used_mem_slots=row['used_mem_slots'],
            used_cpu_slots=row['used_cpu_slots'],
            used_gpu_slots=row['used_gpu_slots'],
            used_tpu_slots=row['used_tpu_slots'],
            addr=row['addr'],
            first_contact=row['first_contact'],
            lost_at=row['lost_at'],
        )

    async def resolve_computations(self, info, status=None):
        '''
        Retrieves all children worker sessions run by this agent.
        '''
        manager = info.context['dlmgr']
        loader = manager.get_loader('Computation.by_agent_id', status=status)
        return await loader.load(self.id)

    @staticmethod
    async def load_all(dbpool, status=None):
        async with dbpool.acquire() as conn:
            query = (sa.select('*')
                       .select_from(agents))
            if status is not None:
                status = AgentStatus[status]
                query = query.where(agents.c.status == status)
            result = await conn.execute(query)
            rows = await result.fetchall()
            return [Agent.from_row(r) for r in rows]

    @staticmethod
    async def batch_load(dbpool, agent_ids, status=None):
        async with dbpool.acquire() as conn:
            query = (sa.select('*')
                       .select_from(agents)
                       .where(agents.c.id.in_(agent_ids))
                       .order_by(agents.c.id))
            if status is not None:
                status = AgentStatus[status]
                query = query.where(agents.c.status == status)
            objs_per_key = OrderedDict()
            for k in agent_ids:
                objs_per_key[k] = None
            async for row in conn.execute(query):
                o = Agent.from_row(row)
                objs_per_key[row.id] = o
        return tuple(objs_per_key.values())
