# Copyright 2022 The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Utility functions for tests."""

import contextlib
import copy
from collections.abc import Iterator
from datetime import datetime, timedelta, timezone
from itertools import count
from typing import Any

try:
    import pydantic.v1 as pydantic
except ImportError:
    import pydantic as pydantic  # type: ignore

from debusine.client.models import (
    ArtifactResponse,
    PaginatedResponse,
    StrictBaseModel,
    WorkRequestResponse,
    WorkflowTemplateResponse,
)
from debusine.tasks import BaseTask


_data_generator_counter = count(0)


def data_generator(size: int) -> Iterator[bytes]:
    """
    Return a data generator of size bytes.

    Each time generates different data.
    """
    while True:
        chunk = bytearray()

        while len(chunk) < size:
            s = str(next(_data_generator_counter)).encode("ascii")
            remaining_space = size - len(chunk)
            chunk.extend(s[:remaining_space])
        yield bytes(chunk)


def yesterday() -> datetime:
    """Return datetime of yesterday."""
    return datetime.now(timezone.utc) - timedelta(days=1)


def tomorrow() -> datetime:
    """Return datetime of tomorrow."""
    return datetime.now(timezone.utc) + timedelta(days=1)


def date_time_to_isoformat_rest_framework(value: datetime) -> str:
    """
    Django REST serialize datetime in a specific way.

    It is reimplemented (simplified) here for Debusine client tests (which
    do not use Django REST).

    See rest_framework/fields.py, DateTimeField.to_representation.
    """
    value_str = value.isoformat()
    if value_str.endswith("+00:00"):
        value_str = value_str[:-6] + "Z"

    return value_str


def create_artifact_response(**kwargs) -> ArtifactResponse:
    """Return an ArtifactResponse. Use defaults for certain fields."""
    defaults: dict[str, Any] = {
        "category": "Testing",
        "data": {},
        "files_to_upload": [],
        "created_at": datetime.utcnow(),
        "workspace": "Testing",
        "download_tar_gz_url": "https://example.com/some-path",
    }
    defaults.update(kwargs)

    return ArtifactResponse(**defaults)


def create_work_request_response(**kwargs) -> WorkRequestResponse:
    """Return a WorkRequestResponse. Use defaults for certain fields."""
    defaults: dict[str, Any] = {
        "id": 11,
        "created_at": datetime.utcnow(),
        "status": "pending",
        "result": "",
        "task_type": "Worker",
        "task_name": "sbuild",
        "task_data": {},
        "priority_base": 0,
        "priority_adjustment": 0,
        "artifacts": [],
        "workspace": "Testing",
    }
    defaults.update(kwargs)

    return WorkRequestResponse(**defaults)


def create_workflow_template_response(**kwargs) -> WorkflowTemplateResponse:
    """Return a WorkflowTemplateResponse. Use defaults for certain fields."""
    defaults: dict[str, Any] = {
        "id": 11,
        "task_name": "noop",
        "task_data": {},
        "priority": 0,
        "workspace": "Testing",
    }
    defaults.update(kwargs)

    return WorkflowTemplateResponse(**defaults)


def create_listing_response(
    *args: StrictBaseModel, next_url: str | None = None
) -> PaginatedResponse:
    """Return a single page of PaginatedResponse, containing args."""
    return PaginatedResponse(
        count=None,
        next=(
            pydantic.parse_obj_as(pydantic.AnyUrl, next_url)
            if next_url is not None
            else None
        ),
        previous=None,
        results=[arg.dict() for arg in args],
    )


@contextlib.contextmanager
def preserve_task_registry():
    """Make task registry changes ephemeral."""
    orig = BaseTask._sub_tasks
    BaseTask._sub_tasks = copy.deepcopy(BaseTask._sub_tasks)
    try:
        yield
    finally:
        BaseTask._sub_tasks = orig
