#
# Copyright (C) 2021 GNS3 Technologies Inc.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
API routes for images.
"""

import os
import logging
import urllib.parse

from fastapi import APIRouter, Request, Depends, status
from fastapi.encoders import jsonable_encoder
from starlette.requests import ClientDisconnect
from sqlalchemy.orm.exc import MultipleResultsFound
from typing import List, Optional

from gns3server import schemas
from gns3server.config import Config
from gns3server.compute.qemu import Qemu
from gns3server.utils.images import (
    InvalidImageError,
    write_image,
    read_image_info,
    default_images_directory,
    get_builtin_disks,
)
from gns3server.db.repositories.images import ImagesRepository
from gns3server.db.repositories.templates import TemplatesRepository
from gns3server.db.repositories.rbac import RbacRepository
from gns3server.controller import Controller
from gns3server.controller.controller_error import (
    ControllerError,
    ControllerNotFoundError,
    ControllerForbiddenError,
    ControllerBadRequestError,
)

from .dependencies.authentication import get_current_active_user
from .dependencies.database import get_repository
from .dependencies.rbac import has_privilege

log = logging.getLogger(__name__)

router = APIRouter()


@router.post(
    "/qemu/{image_path:path}",
    response_model=schemas.Image,
    status_code=status.HTTP_201_CREATED,
    dependencies=[Depends(has_privilege("Image.Allocate"))],
)
async def create_qemu_image(
    image_path: str,
    image_data: schemas.QemuDiskImageCreate,
    images_repo: ImagesRepository = Depends(get_repository(ImagesRepository)),
) -> schemas.Image:
    """
    Create a new blank Qemu image.

    Required privilege: Image.Allocate
    """

    allow_raw_image = Config.instance().settings.Server.allow_raw_images
    if image_data.format == schemas.QemuDiskImageFormat.raw and not allow_raw_image:
        raise ControllerBadRequestError("Raw images are not allowed")

    disk_image_path = urllib.parse.unquote(image_path)
    image_dir, image_name = os.path.split(disk_image_path)
    # check if the path is within the default images directory
    base_images_directory = os.path.expanduser(Config.instance().settings.Server.images_path)
    full_path = os.path.abspath(os.path.join(base_images_directory, image_dir, image_name))
    if os.path.commonprefix([base_images_directory, full_path]) != base_images_directory:
        raise ControllerForbiddenError(f"Cannot write disk image, '{disk_image_path}' is forbidden")

    if not image_dir:
        # put the image in the default images directory for Qemu
        directory = default_images_directory(image_type="qemu")
        os.makedirs(directory, exist_ok=True)
        disk_image_path = os.path.abspath(os.path.join(directory, disk_image_path))

    if await images_repo.get_image(disk_image_path):
        raise ControllerBadRequestError(f"Disk image '{disk_image_path}' already exists")

    options = jsonable_encoder(image_data, exclude_unset=True)
    # FIXME: should we have the create_disk_image in the compute code since
    # this code is used to create images on the controller?
    await Qemu.instance().create_disk_image(disk_image_path, options)

    image_info = await read_image_info(disk_image_path, "qemu")

    image = await images_repo.get_image(disk_image_path)
    if image:
        # the image has already been added to the database
        return image
    else:
        return await images_repo.add_image(**image_info)


@router.get("", response_model=List[schemas.Image], dependencies=[Depends(has_privilege("Image.Audit"))])
async def get_images(
    images_repo: ImagesRepository = Depends(get_repository(ImagesRepository)),
    image_type: Optional[schemas.ImageType] = None,
) -> List[schemas.Image]:
    """
    Return all images.

    Required privilege: Image.Audit
    """

    return await images_repo.get_images(image_type)


@router.post(
    "/upload/{image_path:path}",
    response_model=schemas.Image,
    status_code=status.HTTP_201_CREATED,
    dependencies=[Depends(has_privilege("Image.Allocate"))],
)
async def upload_image(
    image_path: str,
    request: Request,
    images_repo: ImagesRepository = Depends(get_repository(ImagesRepository)),
    templates_repo: TemplatesRepository = Depends(get_repository(TemplatesRepository)),
    current_user: schemas.User = Depends(get_current_active_user),
    rbac_repo: RbacRepository = Depends(get_repository(RbacRepository)),
    install_appliances: Optional[bool] = False,
) -> schemas.Image:
    """
    Upload an image.

    Example: curl -X POST http://host:port/v3/images/upload/my_image_name.qcow2 \
    -H 'Authorization: Bearer <token>' --data-binary @"/path/to/image.qcow2"

    Required privilege: Image.Allocate
    """

    image_path = urllib.parse.unquote(image_path)
    image_dir, image_name = os.path.split(image_path)
    # check if the path is within the default images directory
    base_images_directory = os.path.expanduser(Config.instance().settings.Server.images_path)
    full_path = os.path.abspath(os.path.join(base_images_directory, image_dir, image_name))
    if os.path.commonprefix([base_images_directory, full_path]) != base_images_directory:
        raise ControllerForbiddenError(f"Cannot write image, '{image_path}' is forbidden")

    # If the client sends X-MD5-Checksum, check for a duplicate before consuming the upload stream
    checksum_header = request.headers.get("X-MD5-Checksum")
    if checksum_header:
        check_dir = os.path.dirname(full_path) if image_dir else None
        duplicate = await images_repo.get_image_by_checksum(checksum_header, check_dir)
        if duplicate:
            location = f" in '{check_dir}'" if check_dir else ""
            raise ControllerError(f"Image '{duplicate.filename}' with the same checksum already exists{location}")

    try:
        allow_raw_image = Config.instance().settings.Server.allow_raw_images
        image = await write_image(image_path, full_path, request.stream(), images_repo, allow_raw_image=allow_raw_image)
    except (OSError, InvalidImageError, ClientDisconnect) as e:
        raise ControllerError(f"Could not save image '{image_path}': {e}")

    if install_appliances:
        # attempt to automatically create templates based on image checksum
        await Controller.instance().appliance_manager.install_appliances_from_image(
            image_path,
            image.checksum,
            images_repo,
            templates_repo,
            rbac_repo,
            current_user,
            os.path.dirname(image.path),
        )

    return image


@router.delete(
    "/prune", status_code=status.HTTP_204_NO_CONTENT, dependencies=[Depends(has_privilege("Image.Allocate"))]
)
async def prune_images(
    images_repo: ImagesRepository = Depends(get_repository(ImagesRepository)),
) -> None:
    """
    Prune images not attached to any template.

    Images referenced by a node in any project (opened or closed) are kept.

    Required privilege: Image.Allocate
    """

    skip_images = get_builtin_disks()
    # a single pass over all projects' node properties protects every
    # referenced file name at once
    referenced_filenames = Controller.instance().collect_referenced_image_filenames()
    await images_repo.prune_images(list(skip_images) + list(referenced_filenames))


@router.post("/install", status_code=status.HTTP_200_OK, dependencies=[Depends(has_privilege("Image.Allocate"))])
async def install_images(
    images_repo: ImagesRepository = Depends(get_repository(ImagesRepository)),
    templates_repo: TemplatesRepository = Depends(get_repository(TemplatesRepository)),
) -> dict:
    """
    Attempt to automatically create templates based on image checksums.

    Returns the list of created templates and the list of skipped
    candidates (with the reason why they were skipped).

    Required privilege: Image.Allocate
    """

    created = []
    skipped = []
    skip_images = get_builtin_disks()
    images = await images_repo.get_images()
    for image in images:
        if skip_images and image.filename in skip_images:
            log.debug(f"Skipping image '{image.path}' for image installation")
            continue
        templates = await images_repo.get_image_templates(image.image_id)
        if templates:
            # the image is already used by a template
            log.warning(f"Image '{image.path}' is used by one or more templates")
            skipped.append(
                {
                    "name": image.filename,
                    "reason": "image is already used by one or more templates",
                }
            )
            continue
        results = await Controller.instance().appliance_manager.install_appliances_from_image(
            image.path, image.checksum, images_repo, templates_repo, None, None, os.path.dirname(image.path)
        )
        for result in results:
            if result.get("status") == "created":
                created.append({k: v for k, v in result.items() if k != "status"})
            else:
                skipped.append({k: v for k, v in result.items() if k != "status"})
    return {"created": created, "skipped": skipped}


@router.get("/{image_path:path}", response_model=schemas.Image, dependencies=[Depends(has_privilege("Image.Audit"))])
async def get_image(
    image_path: str,
    images_repo: ImagesRepository = Depends(get_repository(ImagesRepository)),
) -> schemas.Image:
    """
    Return an image.

    Required privilege: Image.Audit
    """

    image_path = urllib.parse.unquote(image_path)
    image = await images_repo.get_image(image_path)
    if not image:
        raise ControllerNotFoundError(f"Image '{image_path}' not found")
    return image


@router.delete(
    "/{image_path:path}",
    status_code=status.HTTP_204_NO_CONTENT,
    dependencies=[Depends(has_privilege("Image.Allocate"))],
)
async def delete_image(
    image_path: str,
    images_repo: ImagesRepository = Depends(get_repository(ImagesRepository)),
) -> None:
    """
    Delete an image.

    Required privilege: Image.Allocate
    """

    image_path = urllib.parse.unquote(image_path)

    try:
        image = await images_repo.get_image(image_path)
    except MultipleResultsFound:
        raise ControllerBadRequestError(
            f"Image '{image_path}' matches multiple images. Please include the absolute path of the image"
        )

    if not image:
        raise ControllerNotFoundError(f"Image '{image_path}' not found")

    templates = await images_repo.get_image_templates(image.image_id)
    if templates:
        template_names = ", ".join([template.name for template in templates])
        raise ControllerError(f"Image '{image_path}' is used by one or more templates: {template_names}")

    project_names = Controller.instance().find_projects_using_image(image.filename)
    if project_names:
        raise ControllerError(f"Image '{image_path}' is used by one or more projects: {', '.join(project_names)}")

    try:
        os.remove(image.path)
    except OSError:
        log.warning(f"Could not delete image file {image.path}")

    success = await images_repo.delete_image(image_path)
    if not success:
        raise ControllerError(f"Image '{image_path}' could not be deleted")