"""
All or portions of this file Copyright (c) Amazon.com, Inc. or its affiliates or
its licensors.

For complete copyright and license terms please see the LICENSE at the root of this
distribution (the "License"). All use of this software is governed by the License,
or, if provided, by the license below or the license accompanying this file. Do not
remove or modify any license notices. This file is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
"""
import azlmbr.atom
import azlmbr.legacy.general as general

from editor_python_test_tools.editor_test_helper import EditorTestHelper

DEFAULT_FRAME_WIDTH = 1920
DEFAULT_FRAME_HEIGHT = 1080
FOLDER_PATH = '@user@/PythonTests/Automated/Screenshots'

helper = EditorTestHelper(log_prefix="Atom_ScreenshotHelper")


class ScreenshotHelper(object):
    """
    A helper to capture screenshots and wait for them.
    """

    def __init__(self, idle_wait_frames_callback, frame_width=DEFAULT_FRAME_WIDTH, frame_height=DEFAULT_FRAME_HEIGHT):
        super().__init__()
        self.done = False
        self.capturedScreenshot = False
        self.max_frames_to_wait = 60
        self.prepare_viewport_for_screenshot(frame_width, frame_height)
        self.idle_wait_frames_callback = idle_wait_frames_callback

    def capture_screenshot_blocking_in_game_mode(self, filename):
        helper.enter_game_mode(["", ""])
        general.idle_wait_frames(120)
        self.capture_screenshot_blocking(filename)
        helper.exit_game_mode(["", ""])

    def prepare_viewport_for_screenshot(self, frame_width, frame_height):
        cur_viewport_size = general.get_viewport_size()
        if int(cur_viewport_size.x) != frame_width or int(cur_viewport_size.y) != frame_height:
            general.set_viewport_expansion_policy('FixedSize')
            general.idle_wait_frames(1)
            general.set_viewport_size(frame_width, frame_height)
            general.set_cvar_integer('r_DisplayInfo', 0)
            general.update_viewport()
            general.idle_wait_frames(120)

        new_viewport_size = general.get_viewport_size()
        if int(new_viewport_size.x) != frame_width or int(new_viewport_size.y) != frame_height:
            general.log("Resolution is incorrect!")
            general.log(f"width: {int(new_viewport_size.x)}")
            general.log(f"height: {int(new_viewport_size.y)}")

    def capture_screenshot_blocking(self, filename, folder_path=FOLDER_PATH):
        """
        Capture a screenshot and block the execution until the screenshot has been written to the disk.
        """
        self.handler = azlmbr.atom.FrameCaptureNotificationBusHandler()
        self.handler.connect()
        self.handler.add_callback('OnCaptureFinished', self.on_screenshot_captured)

        self.done = False
        self.capturedScreenshot = False
        success = azlmbr.atom.FrameCaptureRequestBus(
            azlmbr.bus.Broadcast, "CaptureScreenshot", f"{folder_path}/{filename}")
        if success:
            self.wait_until_screenshot()
            general.log("Screenshot taken.")
        else:
            general.log("Screenshot failed")
        return self.capturedScreenshot

    def on_screenshot_captured(self, parameters):
        # the parameters come in as a tuple
        if parameters[0] == azlmbr.atom.FrameCaptureResult_Success:
            general.log(f"screenshot saved: {parameters[1]}")
            self.capturedScreenshot = True
        else:
            general.log(f"screenshot failed: {parameters[1]}")
        self.done = True
        self.handler.disconnect()

    def wait_until_screenshot(self):
        frames_waited = 0
        while self.done == False:
            self.idle_wait_frames_callback(1)
            if frames_waited > self.max_frames_to_wait:
                general.log("timeout while waiting for the screenshot to be written")
                self.handler.disconnect()
                break
            else:
                frames_waited = frames_waited + 1
        general.log(f"(waited {frames_waited} frames)")