# SPDX-FileCopyrightText: 2024 UL Research Institutes
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

# mypy: disable-error-code="import-untyped"
import time
from datetime import datetime, timedelta, timezone

import pytest
from conftest import (  # type: ignore[import-not-found]
    DATA_DIR,
    wait_for_pipeline_run_status,
    wait_for_pipeline_run_success,
    wait_for_success,
)

from dyff.audit.local.platform import DyffLocalPlatform
from dyff.client import Client
from dyff.schema import commands
from dyff.schema.platform import (
    Challenge,
    ChallengeContent,
    ChallengeContentPage,
    ChallengeRules,
    ChallengeTask,
    ChallengeTaskContent,
    ChallengeTaskExecutionEnvironment,
    ChallengeTaskExecutionEnvironmentChoices,
    ChallengeTaskRules,
    ChallengeTaskSchedule,
    ChallengeTaskVisibility,
    ChallengeVisibility,
    Entities,
    EntityIdentifier,
    InferenceService,
    Method,
    Pipeline,
    QuotaSelector,
    SecurityContext,
    Submission,
    SubmissionStructure,
    Team,
    TeamAffiliation,
    TeamMember,
)
from dyff.schema.requests import (
    ChallengeContentEditRequest,
    ChallengeCreateRequest,
    ChallengeRulesEditRequest,
    ChallengeTaskCreateRequest,
    ChallengeTaskRulesEditRequest,
    ChallengeTeamCreateRequest,
    ResubmissionRequest,
    SubmissionCreateRequest,
    TeamEditRequest,
)


@pytest.mark.datafiles(DATA_DIR)
def test_challenges_create(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    account: str = ctx["account"]

    challenge = dyffapi.challenges.create(
        ChallengeCreateRequest(
            account=account,
            content=ChallengeContent(page=ChallengeContentPage(summary="summary")),
            rules=ChallengeRules(
                visibility=ChallengeVisibility(
                    teams={"*": "submitter", "nobody": "public"}
                )
            ),
        )
    )

    print(f"challenge: {challenge.id}")
    ctx["challenge"] = challenge

    wait_for_success(
        lambda: dyffapi.challenges.get(challenge.id),
        timeout=timedelta(minutes=2),
    )


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_challenges_create",
    ]
)
def test_challenges_edit_content(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    challenge: Challenge = ctx["challenge"]

    dyffapi.challenges.edit_content(
        challenge.id,
        ChallengeContentEditRequest(
            content=commands.ChallengeContentPatch(
                page=commands.ChallengeContentPagePatch(summary="", body="body")
            )
        ),
    )

    time.sleep(10)

    updated = dyffapi.challenges.get(challenge.id)
    assert updated.content.page.title == "Untitled Challenge"
    assert updated.content.page.summary == ""
    assert updated.content.page.body == "body"


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_challenges_create",
        "test_challenges_create_team",
    ]
)
def test_challenges_edit_rules(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    challenge: Challenge = ctx["challenge"]
    team: Team = ctx["team"]

    dyffapi.challenges.edit_rules(
        challenge.id,
        ChallengeRulesEditRequest(
            rules=commands.ChallengeRulesPatch(
                visibility=commands.ChallengeRulesVisibilityPatch(
                    teams={"nobody": None, team.id: "public"}
                )
            )
        ),
    )

    time.sleep(10)

    updated = dyffapi.challenges.get(challenge.id)
    assert updated.rules.visibility.teams == {
        "*": "submitter",
        team.id: "public",
    }


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "tests/test_pipelines.py::test_pipelines_create",
        "test_challenges_create",
    ]
)
def test_challenges_create_task(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    account: str = ctx["account"]
    challenge: Challenge = ctx["challenge"]
    method: Method = ctx["pipeline_method"]
    pipeline: Pipeline = ctx["pipeline"]

    request = ChallengeTaskCreateRequest(
        account=account,
        challenge=challenge.id,
        name="test",
        assessment=pipeline.id,
        rules=ChallengeTaskRules(
            executionEnvironment=ChallengeTaskExecutionEnvironmentChoices(
                choices={
                    "default": ChallengeTaskExecutionEnvironment(
                        cpu="1",
                        memory="1Gi",
                    ),
                    "big": ChallengeTaskExecutionEnvironment(
                        cpu="2",
                        memory="1Gi",
                    ),
                },
            ),
            schedule=ChallengeTaskSchedule(
                openingTime=datetime(year=2000, month=1, day=1, tzinfo=timezone.utc),
            ),
            visibility=ChallengeTaskVisibility(
                scores={
                    method.id: {
                        "float_unit_primary": "public",
                        "int": "public",
                        "*": "reviewer",
                    },
                    "*": {"int_percent": "public"},
                },
            ),
        ),
        content=ChallengeTaskContent(page=ChallengeContentPage(summary="summary")),
        submissionStructure=SubmissionStructure(
            submissionKind=Entities.InferenceService.value,
            pipelineKeyword="submission",
        ),
    )
    task = dyffapi.challenges.create_task(challenge.id, request)

    print(f"challengetask: {task.id}")
    ctx["challengetask"] = task

    time.sleep(10)
    stored_task = dyffapi.challenges.get(challenge.id).tasks[task.id]
    assert stored_task.model_dump() == task.model_dump()


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_challenges_create_task",
    ]
)
def test_challenges_edit_task_content(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    challenge: Challenge = ctx["challenge"]
    task: ChallengeTask = ctx["challengetask"]

    dyffapi.challenges.edit_task_content(
        challenge.id,
        task.id,
        ChallengeContentEditRequest(
            content=commands.ChallengeContentPatch(
                page=commands.ChallengeContentPagePatch(summary="", body="body")
            )
        ),
    )

    time.sleep(10)

    updated = dyffapi.challenges.get(challenge.id)
    updated_task = updated.tasks[task.id]
    assert updated_task.content.page.title == "Untitled Challenge"
    assert updated_task.content.page.summary == ""
    assert updated_task.content.page.body == "body"


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_challenges_create_task",
    ]
)
def test_challenges_edit_task_rules(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    challenge: Challenge = ctx["challenge"]
    task: ChallengeTask = ctx["challengetask"]
    method: Method = ctx["pipeline_method"]
    pipeline: Pipeline = ctx["pipeline"]

    dyffapi.challenges.edit_task_rules(
        challenge.id,
        task.id,
        ChallengeTaskRulesEditRequest.model_validate(
            {
                "rules": {
                    "executionEnvironment": {
                        "choices": {
                            "default": {
                                "cpu": "2",
                                "memory": "1Gi",
                            },
                            "big": None,
                            "small": {
                                "cpu": "1",
                                "memory": "512Mi",
                            },
                        }
                    },
                    "schedule": {
                        "openingTime": None,
                        "closingTime": datetime(
                            year=2032, month=1, day=1, tzinfo=timezone.utc
                        ),
                        "submissionLimitPerCycle": 1,
                    },
                    "visibility": {
                        "scores": {
                            method.id: {
                                "float_unit_primary": None,
                                "int": "reviewer",
                            },
                            "*": None,
                        },
                    },
                },
                "assessment": "test",
            }
        ),
    )

    time.sleep(10)

    updated = dyffapi.challenges.get(challenge.id)
    updated_task = updated.tasks[task.id]
    assert updated_task.rules.executionEnvironment.model_dump() == {
        "choices": {
            "default": {
                "accelerators": {},
                "cpu": "2",
                "memory": "1Gi",
            },
            "small": {
                "accelerators": {},
                "cpu": "1",
                "memory": "512Mi",
            },
        }
    }
    # The round-trip encoding results in a different type of tzinfo object
    # even though the two are semantically equivalent
    assert updated_task.rules.schedule.model_dump(mode="json") == ChallengeTaskSchedule(
        closingTime=datetime(year=2032, month=1, day=1, tzinfo=timezone.utc),
        submissionLimitPerCycle=1,
    ).model_dump(mode="json")

    assert updated_task.rules.visibility.model_dump() == {
        "scores": {
            method.id: {
                "int": "reviewer",
                "*": "reviewer",
            }
        }
    }

    assert updated_task.assessment == "test"

    # We have to change the pipeline back so that the 'submit' test works
    dyffapi.challenges.edit_task_rules(
        challenge.id,
        task.id,
        ChallengeTaskRulesEditRequest.model_validate(
            {
                "assessment": pipeline.id,
            }
        ),
    )
    time.sleep(10)

    reverted = dyffapi.challenges.get(challenge.id)
    reverted_task = reverted.tasks[task.id]
    assert reverted_task.assessment == pipeline.id


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_challenges_create",
    ]
)
def test_challenges_create_team(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    account: str = ctx["account"]
    challenge: Challenge = ctx["challenge"]
    request = ChallengeTeamCreateRequest(
        account=account,
        members={
            "rusty": TeamMember(
                name="Rusty Shackleford",
                isCorrespondingMember=True,
                affiliations=["dales_dead_bug"],
            ),
            "peggy": TeamMember(
                name="Peggy Hill",
                isCorrespondingMember=False,
                affiliations=["tom_landry_middle_school"],
            ),
        },
        affiliations={
            "dales_dead_bug": TeamAffiliation(
                name="Dale's Dead Bug",
            ),
            "tom_landry_middle_school": TeamAffiliation(
                name="Tom Landry Middle School",
            ),
        },
    )
    team = dyffapi.challenges.create_team(challenge.id, request)

    print(f"team: {team.id}")
    ctx["team"] = team

    time.sleep(10)
    stored_team = dyffapi.teams.get(team.id)
    assert stored_team.members == team.members
    assert stored_team.affiliations == team.affiliations


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_challenges_create_team",
    ]
)
def test_challenges_list_teams(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    challenge: Challenge = ctx["challenge"]

    teams = dyffapi.challenges.teams(challenge.id)
    assert len(teams) == 1


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_challenges_create_team",
    ]
)
def test_challenges_edit_team(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    team: Team = ctx["team"]

    dyffapi.teams.edit(
        team.id,
        TeamEditRequest(
            members={
                "rusty": None,
                "hank": TeamMember(
                    name="Hank Hill",
                    isCorrespondingMember=True,
                    affiliations=["strickland_propane"],
                ),
            },
            affiliations={
                "dales_dead_bug": None,
                "strickland_propane": TeamAffiliation(
                    name="Strickland Propane",
                    note="Taste the meat, not the heat!",
                ),
            },
        ),
    )

    time.sleep(10)
    stored_team = dyffapi.teams.get(team.id)
    assert stored_team.members == {
        "peggy": TeamMember(
            name="Peggy Hill",
            isCorrespondingMember=False,
            affiliations=["tom_landry_middle_school"],
        ),
        "hank": TeamMember(
            name="Hank Hill",
            isCorrespondingMember=True,
            affiliations=["strickland_propane"],
        ),
    }
    assert stored_team.affiliations == {
        "tom_landry_middle_school": TeamAffiliation(name="Tom Landry Middle School"),
        "strickland_propane": TeamAffiliation(
            name="Strickland Propane",
            note="Taste the meat, not the heat!",
        ),
    }


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        # The real value of the .assessment field gets set in the
        # 'edit_task_rules' test; otherwise there's an ordering ambiguity
        "test_challenges_edit_task_rules",
    ]
)
def test_challenges_submit(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if pytestconfig.getoption("skip_challenges_submit"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    account: str = ctx["account_non_exempt"]
    challenge: Challenge = ctx["challenge"]
    task: ChallengeTask = ctx["challengetask"]
    team: Team = ctx["team"]
    inferenceService: InferenceService = ctx["inferenceservice_mock"]

    submission = dyffapi.challenges.submit(
        challenge.id,
        task.id,
        SubmissionCreateRequest(
            account=account,
            team=team.id,
            submission=EntityIdentifier.of(inferenceService),
        ),
    )

    print(f"challengesubmission {submission.id}")
    ctx["challengesubmission"] = submission

    wait_for_pipeline_run_success(
        lambda: dyffapi.pipelines.get_run_status(submission.pipelineRun),
        timeout=timedelta(minutes=10),
    )


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        # The real value of the .assessment field gets set in the
        # 'edit_task_rules' test; otherwise there's an ordering ambiguity
        "test_challenges_submit",
    ]
)
def test_challenges_rate_quota(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if pytestconfig.getoption("skip_challenges_submit"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    account: str = ctx["account_non_exempt"]
    challenge: Challenge = ctx["challenge"]
    task: ChallengeTask = ctx["challengetask"]
    team: Team = ctx["team"]
    inferenceService: InferenceService = ctx["inferenceservice_mock"]

    submission = dyffapi.challenges.submit(
        challenge.id,
        task.id,
        SubmissionCreateRequest(
            account=account,
            team=team.id,
            submission=EntityIdentifier.of(inferenceService),
        ),
    )

    print(f"challengesubmission_rate_quota {submission.id}")
    ctx["challengesubmission_rate_quota"] = submission

    time.sleep(10)

    submission = dyffapi.submissions.get(submission.id)
    # TODO: This test currently fails (DYFF-959)
    # assert submission.status == "Error"
    assert submission.reason == "RateLimitExceeded"

    pipeline_status = dyffapi.pipelines.get_run_status(submission.pipelineRun)
    assert pipeline_status.nodes["evaluation"].status == "Error"


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_challenges_submit",
    ]
)
def test_challenges_resubmit(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if pytestconfig.getoption("skip_challenges_submit"):
        pytest.skip()
    if not pytestconfig.getoption("enable_extra"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    account: str = ctx["account_non_exempt"]
    challenge: Challenge = ctx["challenge"]
    task: ChallengeTask = ctx["challengetask"]
    submission: Submission = ctx["challengesubmission"]

    # Re-submit twice to test disabling rate quota
    assert task.rules.schedule.submissionLimitPerCycle == 1

    resubmission1 = dyffapi.submissions.resubmit(
        submission.id,
        ResubmissionRequest(
            account=account,
            securityContext=SecurityContext(
                trusted=False,
                ignoreQuotas=[
                    QuotaSelector(principal="*", resource="*", flavor="Rate")
                ],
            ),
        ),
    )
    assert resubmission1.resubmissionOf == submission.id

    print(f"resubmission1 {resubmission1.id}")
    ctx["resubmission1"] = resubmission1

    resubmission2 = dyffapi.submissions.resubmit(
        submission.id,
        ResubmissionRequest(
            account=account,
            securityContext=SecurityContext(
                trusted=False,
                ignoreQuotas=[
                    QuotaSelector(principal="*", resource="*", flavor="Rate")
                ],
            ),
        ),
    )
    assert resubmission2.resubmissionOf == submission.id

    print(f"resubmission2 {resubmission2.id}")
    ctx["resubmission2"] = resubmission2

    wait_for_pipeline_run_success(
        lambda: dyffapi.pipelines.get_run_status(resubmission1.pipelineRun),
        timeout=timedelta(minutes=10),
    )

    wait_for_pipeline_run_success(
        lambda: dyffapi.pipelines.get_run_status(resubmission2.pipelineRun),
        timeout=timedelta(minutes=10),
    )


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        # The real value of the .assessment field gets set in the
        # 'edit_task_rules' test; otherwise there's an ordering ambiguity
        "test_challenges_edit_task_rules",
    ]
)
def test_challenges_terminate_submission(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_challenges"):
        pytest.skip()
    if pytestconfig.getoption("skip_challenges_submit"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    account: str = ctx["account"]
    challenge: Challenge = ctx["challenge"]
    task: ChallengeTask = ctx["challengetask"]
    team: Team = ctx["team"]
    inferenceService: InferenceService = ctx["inferenceservice_mock"]

    submission = dyffapi.challenges.submit(
        challenge.id,
        task.id,
        SubmissionCreateRequest(
            account=account,
            team=team.id,
            submission=EntityIdentifier.of(inferenceService),
        ),
    )

    print(f"challengesubmission_terminate {submission.id}")

    wait_for_success(
        lambda: dyffapi.challenges.get(challenge.id),
        timeout=timedelta(minutes=2),
    )

    # Wait until a step is actually running
    wait_for_pipeline_run_status(
        lambda: dyffapi.pipelines.get_run_status(submission.pipelineRun),
        timeout=timedelta(minutes=10),
        success={
            "evaluation": ["Admitted", "Complete"],
            "safetycase": ["Admitted", "Complete"],
        },
        failure={
            "evaluation": ["Error", "Failed"],
            "safetycase": ["Error", "Failed"],
        },
        polling_interval_seconds=5,
    )

    dyffapi.submissions.terminate(submission.id)
    time.sleep(10)

    submission = dyffapi.submissions.get(submission.id)
    pipeline_status = dyffapi.pipelines.get_run_status(submission.pipelineRun)
    assert submission.status == "Terminated"
    assert pipeline_status.status == "Terminated"
    for node, status in pipeline_status.nodes.items():
        # Downstream workflows error when upstream workflows terminate
        # Downstream workflows that already completed don't change status
        assert (
            status.status == "Terminated"
            or (status.status == "Error" and status.reason == "FailedDependency")
            or status.status == "Complete"
        )
