#
# Copyright (c) 2021 the Hugging Face team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.#
from cryptography.hazmat.primitives.asymmetric import rsa
from fastapi import APIRouter, Body, Depends, HTTPException, Path
from starlette.status import HTTP_201_CREATED, HTTP_401_UNAUTHORIZED, HTTP_404_NOT_FOUND

from app.api.dependencies import crypto
from app.api.dependencies.database import get_repository
from app.db.repositories.experiments import ExperimentsRepository
from app.models.experiment import (
    ExperimentCreate,
    ExperimentCreatePublic,
    ExperimentInDB,
    ExperimentPublic,
    ExperimentUpdate,
)
from app.models.experiment_join import ExperimentJoinInput, ExperimentJoinOutput
from app.services.authentication import MoonlandingUser, RepoRole, authenticate


router = APIRouter()


@router.post("/", response_model=ExperimentPublic, name="experiments:create-experiment", status_code=HTTP_201_CREATED)
async def create_new_experiment(
    new_experiment: ExperimentCreatePublic = Body(..., embed=True),
    experiments_repo: ExperimentsRepository = Depends(get_repository(ExperimentsRepository)),
    user: MoonlandingUser = Depends(authenticate),
) -> ExperimentPublic:
    # collaborators_list = new_experiment.collaborators

    if new_experiment.organization_name not in [org.name for org in user.orgs if org.role_in_org == RepoRole.admin]:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail=f"You need to be an admin of the organization {new_experiment.organization_name} to create a collaborative experiment for the model {new_experiment.model_name}",
        )

    experiment = await experiments_repo.get_experiment_by_organization_and_model_name(
        organization_name=new_experiment.organization_name, model_name=new_experiment.model_name
    )

    if experiment is not None:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail=f"An experiment already exist for the organization {new_experiment.organization_name} and the model {new_experiment.model_name}",
        )

    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    public_key = private_key.public_key()

    new_experiment = new_experiment.dict()

    new_experiment_item = ExperimentCreate(
        **new_experiment,
        auth_server_private_key=crypto.save_private_key(private_key),
        auth_server_public_key=crypto.save_public_key(public_key),
    )
    created_experiment_item = await experiments_repo.create_experiment(
        new_experiment=new_experiment_item, requesting_user=user
    )

    created_experiment_item = await get_experiment_by_id(
        id=created_experiment_item.id,
        experiments_repo=experiments_repo,
        user=user,
    )
    return created_experiment_item


@router.get("/", response_model=ExperimentPublic, name="experiments:get-experiment-by-organization-and-model-name")
async def get_experiment_by_organization_and_model_name(
    organization_name: str,
    model_name: str,
    experiments_repo: ExperimentsRepository = Depends(get_repository(ExperimentsRepository)),
    user: MoonlandingUser = Depends(authenticate),
) -> ExperimentPublic:
    experiment = await experiments_repo.get_experiment_by_organization_and_model_name(
        organization_name=organization_name, model_name=model_name
    )

    if not experiment:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail="You need to be an admin of the organization to get the collaborative experiment for the model",
        )

    if experiment.organization_name not in [org.name for org in user.orgs if org.role_in_org == RepoRole.admin]:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail=f"You need to be an admin of the organization {experiment.organization_name} to get the collaborative experiment for the model {experiment.model_name}",
        )

    experiment_public = ExperimentPublic(**experiment.dict())
    return experiment_public


@router.get("/{id}/", response_model=ExperimentPublic, name="experiments:get-experiment-by-id")
async def get_experiment_by_id(
    id: int,
    experiments_repo: ExperimentsRepository = Depends(get_repository(ExperimentsRepository)),
    user: MoonlandingUser = Depends(authenticate),
) -> ExperimentPublic:

    experiment = await experiments_repo.get_experiment_by_id(id=id)

    if not experiment:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail="You need to be an admin of the organization to get the collaborative experiment for the model",
        )

    if experiment.organization_name not in [org.name for org in user.orgs if org.role_in_org == RepoRole.admin]:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail=f"You need to be an admin of the organization {experiment.organization_name} to get the collaborative experiment for the model {experiment.model_name}",
        )

    experiment_public = ExperimentPublic(**experiment.dict())
    return experiment_public


@router.put("/{id}/", response_model=ExperimentPublic, name="experiments:update-experiment-by-id")
async def update_experiment_by_id(
    id: int = Path(..., ge=1, title="The ID of the experiment to update."),
    experiment_update: ExperimentUpdate = Body(..., embed=True),
    experiments_repo: ExperimentsRepository = Depends(get_repository(ExperimentsRepository)),
    user: MoonlandingUser = Depends(authenticate),
) -> ExperimentPublic:
    experiment = await experiments_repo.get_experiment_by_id(id=id)

    if not experiment:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail="You need to be an admin of the organization to update the collaborative experiment for the model",
        )

    if experiment.organization_name not in [org.name for org in user.orgs if org.role_in_org == RepoRole.admin]:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail=f"You need to be an admin of the organization {experiment.organization_name} to update the collaborative experiment for the model {experiment.model_name}",
        )

    if not experiment:
        raise HTTPException(status_code=HTTP_404_NOT_FOUND, detail="No experiment found with that id.")

    updated_public_experiment = await update_experiment(experiment, experiment_update, user, experiments_repo)
    return updated_public_experiment


async def update_experiment(
    experiment: ExperimentInDB,
    experiment_update: ExperimentUpdate,
    user: MoonlandingUser,
    experiments_repo: ExperimentsRepository,
):
    experiment_update = ExperimentUpdate(**experiment_update.dict(exclude_unset=True))
    experiment = await experiments_repo.update_experiment_by_id(
        id_exp=experiment.id, experiment_update=experiment_update
    )

    updated_public_experiment = await get_experiment_by_id(
        id=experiment.id,
        experiments_repo=experiments_repo,
        user=user,
    )
    return updated_public_experiment


@router.delete("/{id}/", response_model=ExperimentPublic, name="experiments:delete-experiment-by-id")
async def delete_experiment_by_id(
    id: int = Path(..., ge=1, title="The ID of the experiment to delete."),
    experiments_repo: ExperimentsRepository = Depends(get_repository(ExperimentsRepository)),
    user: MoonlandingUser = Depends(authenticate),
) -> ExperimentPublic:
    experiment = await experiments_repo.get_experiment_by_id(id=id)

    if not experiment:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail="You need to be an admin of the organization to update the collaborative experiment for the model",
        )

    if experiment.organization_name not in [org.name for org in user.orgs if org.role_in_org == RepoRole.admin]:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail=f"You need to be an admin of the organization {experiment.organization_name} to delete the collaborative experiment for the model {experiment.model_name}",
        )

    deleted_public_experiment = await delete_experiment(experiment, user, experiments_repo)
    return deleted_public_experiment


async def delete_experiment(
    experiment: ExperimentInDB, user: MoonlandingUser, experiments_repo: ExperimentsRepository
):
    deleted_id = await experiments_repo.delete_experiment_by_id(id=experiment.id)
    if not deleted_id:
        raise HTTPException(
            status_code=HTTP_404_NOT_FOUND,
            detail=f"No experiment found with the id {experiment.id} for the collaborative experiment of the model {experiment.model_name} of the organization {experiment.organization_name}.",
        )

    deleted_exp = ExperimentPublic(**experiment.dict())
    return deleted_exp


@router.put(
    "/join",
    response_model=ExperimentJoinOutput,
    name="experiments:join-experiment-by-organization-and-model-name",
)
async def join_experiment_by_organization_and_model_name(
    organization_name: str,
    model_name: str,
    experiment_join_input: ExperimentJoinInput = Body(..., embed=True),
    experiments_repo: ExperimentsRepository = Depends(get_repository(ExperimentsRepository)),
    user: MoonlandingUser = Depends(authenticate),
) -> ExperimentJoinOutput:
    experiment = await experiments_repo.get_experiment_by_organization_and_model_name(
        organization_name=organization_name, model_name=model_name
    )

    if not experiment:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail="You need to be at least a reader of the organization to join the collaborative experiment for the model",
        )

    if experiment.organization_name not in [org.name for org in user.orgs]:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail="Access to the experiment denied.",
        )

    exp_pass = await join_experiment(experiment, user, experiment_join_input)
    return exp_pass


@router.put("/join/{id}/", response_model=ExperimentJoinOutput, name="experiments:join-experiment-by-id")
async def join_experiment_by_id(
    id: int = Path(..., ge=1, title="The ID of the experiment the user wants to join."),
    experiment_join_input: ExperimentJoinInput = Body(..., embed=True),
    experiments_repo: ExperimentsRepository = Depends(get_repository(ExperimentsRepository)),
    user: MoonlandingUser = Depends(authenticate),
) -> ExperimentJoinOutput:
    experiment = await experiments_repo.get_experiment_by_id(id=id)

    if not experiment:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail="You need to be at least a reader of the organization to join the collaborative experiment for the model",
        )

    if experiment.organization_name not in [org.name for org in user.orgs]:
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail="Access to the experiment denied.",
        )

    exp_pass = await join_experiment(experiment, user, experiment_join_input)
    return exp_pass


async def join_experiment(
    experiment: ExperimentInDB, user: MoonlandingUser, experiment_join_input: ExperimentJoinInput
):
    if not experiment:
        raise HTTPException(status_code=HTTP_404_NOT_FOUND, detail="No experiment found with that id.")

    hivemind_access = crypto.create_hivemind_access(
        peer_public_key=experiment_join_input.peer_public_key,
        auth_server_private_key=experiment.auth_server_private_key,
        username=user.username,
    )
    exp_pass = ExperimentJoinOutput(**experiment.dict(), hivemind_access=hivemind_access)
    return exp_pass
