from datetime import datetime
from typing import AsyncIterator

import pytz
from dateutil.parser import isoparse
from smb.common.http_client import BaseHttpClientException, HttpClient


class YangClientException(BaseHttpClientException):
    pass


class NaiveDateTime(YangClientException):
    pass


class YangClient(HttpClient):
    _pool_id: str
    _auth_token: str

    def __init__(self, *, url: str, token: str, pool_id: str):
        super().__init__(url=url)
        self._pool_id = pool_id
        self._auth_token = token

    async def create_task_suite(self, input_values: dict) -> dict:
        task_params = await self.request(
            method="POST",
            uri="api/v1/task-suites",
            expected_statuses=[201],
            headers={"Authorization": f"OAuth {self._auth_token}"},
            json={
                "pool_id": self._pool_id,
                "overlap": 1,
                "tasks": [{"input_values": input_values}],
            },
        )

        return dict(
            id=task_params["id"],
            created_at=pytz.utc.localize(isoparse(task_params["created"])),
        )

    async def list_accepted_assignments(
        self, min_created: datetime
    ) -> AsyncIterator[dict]:
        if not min_created.tzinfo:
            raise NaiveDateTime(
                f"min_created must be aware with timezone: {min_created}."
            )
        utc_min_created = min_created.astimezone(pytz.utc)

        params = {
            "pool_id": self._pool_id,
            "status": "ACCEPTED",
            "limit": "500",
            "sort": "id",
            "created_gte": utc_min_created.replace(tzinfo=None).isoformat(),
        }

        while True:
            chunk = await self.request(
                method="GET",
                uri="api/v1/assignments",
                expected_statuses=[200],
                headers={"Authorization": f"OAuth {self._auth_token}"},
                params=params,
            )

            for assignment in chunk["items"]:
                yield self._compose_assignment_data(assignment)

            if not chunk["has_more"]:
                break

            last_item = chunk["items"][-1]
            params["id_gt"] = last_item["id"]

    def _compose_assignment_data(self, assignment: dict) -> dict:
        return dict(
            id=assignment["id"],
            task_suite_id=assignment["task_suite_id"],
            task=assignment["tasks"][0],
            solution=assignment["solutions"][0]["output_values"],
            accepted_at=pytz.utc.localize(isoparse(assignment["accepted"])),
        )
