#
# 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 ipaddress import IPv4Address, IPv6Address
from typing import List

from fastapi import HTTPException
from starlette.status import HTTP_400_BAD_REQUEST

from app.db.repositories.base import BaseRepository
from app.models.experiment import ExperimentCreate, ExperimentInDB, ExperimentUpdate
from app.services.authentication import MoonlandingUser


COLUMNS = "id, organization_name, model_name, creator, coordinator_ip, coordinator_port, auth_server_public_key, auth_server_private_key, created_at, updated_at"
CREATE_EXPERIMENT_QUERY = """
    INSERT INTO experiments (organization_name, model_name, creator, coordinator_ip, coordinator_port, auth_server_public_key, auth_server_private_key)
    VALUES (:organization_name, :model_name, :creator, :coordinator_ip, :coordinator_port, :auth_server_public_key, :auth_server_private_key)
    RETURNING id, organization_name, model_name, creator, coordinator_ip, coordinator_port, auth_server_public_key, auth_server_private_key, created_at, updated_at;
"""
GET_EXPERIMENT_BY_ID_QUERY = """
    SELECT id, organization_name, model_name, creator, coordinator_ip, coordinator_port, auth_server_public_key, auth_server_private_key, created_at, updated_at
    FROM experiments
    WHERE id = :id;
"""
GET_EXPERIMENT_BY_ORGANIZATON_AND_MODEL_NAME_QUERY = """
    SELECT id, organization_name, model_name, creator, coordinator_ip, coordinator_port, auth_server_public_key, auth_server_private_key, created_at, updated_at
    FROM experiments
    WHERE model_name = :model_name
    AND organization_name = :organization_name;
"""
LIST_ALL_USER_EXPERIMENTS_QUERY = """
    SELECT id, organization_name, model_name, creator, coordinator_ip, coordinator_port, auth_server_public_key, auth_server_private_key, created_at, updated_at
    FROM experiments
    WHERE creator = :creator;
"""
UPDATE_EXPERIMENT_BY_ID_QUERY = """
    UPDATE experiments
    SET organization_name = :organization_name,
        model_name        = :model_name,
        coordinator_ip    = :coordinator_ip,
        coordinator_port  = :coordinator_port,
        creator           = :creator
    WHERE id = :id
    RETURNING id, organization_name, model_name, creator, coordinator_ip, coordinator_port, auth_server_public_key, auth_server_private_key, created_at, updated_at;
"""
DELETE_EXPERIMENT_BY_ID_QUERY = """
    DELETE FROM experiments
    WHERE id = :id
    RETURNING id;
"""


class ExperimentsRepository(BaseRepository):
    """ "
    All database actions associated with the Experiment resource
    """

    async def create_experiment(
        self, *, new_experiment: ExperimentCreate, requesting_user: MoonlandingUser
    ) -> ExperimentInDB:
        new_experiment_table = {**new_experiment.dict(), "creator": requesting_user.username}
        if "coordinator_ip" in new_experiment_table.keys() and (
            isinstance(new_experiment_table["coordinator_ip"], IPv4Address)
            or isinstance(new_experiment_table["coordinator_ip"], IPv6Address)
        ):
            new_experiment_table["coordinator_ip"] = str(new_experiment_table["coordinator_ip"])

        experiment = await self.db.fetch_one(query=CREATE_EXPERIMENT_QUERY, values=new_experiment_table)
        return ExperimentInDB(**experiment)

    async def get_experiment_by_organization_and_model_name(
        self, *, organization_name: str, model_name: str
    ) -> ExperimentInDB:
        experiment = await self.db.fetch_one(
            query=GET_EXPERIMENT_BY_ORGANIZATON_AND_MODEL_NAME_QUERY,
            values={"organization_name": organization_name, "model_name": model_name},
        )
        if not experiment:
            return None

        return ExperimentInDB(**experiment)

    async def get_experiment_by_id(self, *, id: int) -> ExperimentInDB:
        experiment = await self.db.fetch_one(query=GET_EXPERIMENT_BY_ID_QUERY, values={"id": id})
        if not experiment:
            return None

        return ExperimentInDB(**experiment)

    async def list_all_user_experiments(self, requesting_user: MoonlandingUser) -> List[ExperimentInDB]:
        experiment_records = await self.db.fetch_all(
            query=LIST_ALL_USER_EXPERIMENTS_QUERY, values={"creator": requesting_user.username}
        )
        return [ExperimentInDB(**exp) for exp in experiment_records]

    async def update_experiment_by_id(self, *, id_exp: int, experiment_update: ExperimentUpdate) -> ExperimentInDB:
        experiment = await self.get_experiment_by_id(id=id_exp)
        if not experiment:
            return None
        experiment_update_params = experiment.copy(update=experiment_update.dict(exclude_unset=True))
        values = {
            **experiment_update_params.dict(
                exclude={
                    "auth_server_public_key",
                    "auth_server_private_key",
                    "created_at",
                    "updated_at",
                }
            )
        }
        if "coordinator_ip" in values.keys() and (
            isinstance(values["coordinator_ip"], IPv4Address) or isinstance(values["coordinator_ip"], IPv6Address)
        ):
            values["coordinator_ip"] = str(values["coordinator_ip"])
        try:
            updated_experiment = await self.db.fetch_one(query=UPDATE_EXPERIMENT_BY_ID_QUERY, values=values)
        except Exception as e:
            print(e)
            raise HTTPException(status_code=HTTP_400_BAD_REQUEST, detail="Invalid update params.")
        return ExperimentInDB(**updated_experiment)

    async def delete_experiment_by_id(self, *, id: int) -> int:
        experiment = await self.get_experiment_by_id(id=id)
        if not experiment:
            return None
        deleted_id = await self.db.execute(query=DELETE_EXPERIMENT_BY_ID_QUERY, values={"id": id})
        return deleted_id
