import carb
import asyncio
import shutil
import sys
import os
import subprocess
import tempfile
from time import time
import omni.kit.test


async def run_process(args, timeout, capture_output=False):
    returncode = 0
    fail_messages = []
    proc = None
    stdout = None
    try:
        cmd = " ".join(args)
        print(f">>> running process: {cmd}\n")

        async def run_proc():
            nonlocal proc
            proc = await asyncio.create_subprocess_exec(
                *args,
                stdout=subprocess.PIPE if capture_output else None,
                stderr=subprocess.STDOUT,
                bufsize=0,
            )
            await proc.wait()

            nonlocal returncode, stdout
            returncode = proc.returncode
            if capture_output:
                stdout = []
                async for line in proc.stdout:
                    stdout.append(line.decode(errors="replace").replace("\r\n", "\n").replace("\r", "\n"))
            proc = None

        await asyncio.wait_for(run_proc(), timeout)

    except subprocess.CalledProcessError as e:
        returncode = e.returncode
        fail_messages.append(f"subprocess.CalledProcessError was raised: {e.output}")
    except asyncio.TimeoutError:
        fail_messages.append(f"Process timed out (timeout: {timeout}). Terminating.")
        if proc:
            proc.terminate()

    # return code failure check
    if returncode != 0:
        fail_messages.append(f"Process return code: {returncode} != 0.")

    return (returncode, fail_messages, stdout)


class KitLaunchTest(omni.kit.test.AsyncTestCase):
    def __init__(self, tests=()):
        super().__init__(tests)

    async def _run_kit(self, args, portable=True, timeout=180, **kwargs):
        launch_args = [sys.argv[0]] + args
        launch_args += ["--/log/flushStandardStreamOutput=1", "--/app/fastShutdown=1", "--/crashreporter/gatherUserStory=0"]
        if portable:
            launch_args += ["--portable"]
        return await run_process(launch_args, timeout, **kwargs)

    async def _run_kit_and_check(self, args, portable=True, timeout=180):
        code, err_messages, stdout = await self._run_kit(args, portable, timeout)
        if len(err_messages) > 0:
            print(err_messages)
        self.assertEqual(len(err_messages), 0, msg=err_messages)
        self.assertEqual(code, 0)


    async def _run_kit_and_capture_stdout(self, args, portable=True, timeout=180, expect_errors=[], expect_success=True):
        code, err_messages, stdout = await self._run_kit(args, portable, timeout, capture_output=True)
        left_errors = list(expect_errors)
        # Check that all expected errors are present in the output
        # Replace [Error] with [SkippedError] to avoid main test failure
        for line in stdout:
            if "[Error]" in line:
                line = line.replace("[Error]", "[SkippedError]")
                for err in left_errors:
                    if err in line:
                        left_errors.remove(err)
            sys.stdout.write(line)
        self.assertTrue(len(left_errors) == 0, f"Expected errors not found: {left_errors}")

        if expect_success:
            self.assertEqual(code, 0)
        else:
            self.assertNotEqual(code, 0)


    async def _run_kit_with_script(
        self,
        script_content,
        args=[],
        kit_file=None,
        add_exec_check=True,
        portable=True,
        kit_file_name=None,
        script_file=None,
        script_extra_args="",
        timeout=180,
    ):
        ts = int(time())

        # Touch special file and check after running kit that file exist. This way we can communicate failure to execute a script.
        exec_check_file = f"{self._temp_folder}/exec_ok_{ts}"
        if add_exec_check:
            script_content += "\n"
            script_content += "from pathlib import Path\n"
            script_content += f"Path(r'{exec_check_file}').touch()\n"

        if not script_file:
            script_file = f"{self._temp_folder}/script_{ts}.py"
        with open(script_file.strip('"'), "w") as f:
            f.write(script_content)
        carb.log_info(f"dump '{script_file}':\n{script_content}")

        launch_args = []

        # Optionally build kit file
        if kit_file:
            if not kit_file_name:
                kit_file_name = f"test_{ts}.kit"
            kit_file_path = f"{self._temp_folder}/{kit_file_name}"
            with open(kit_file_path, "w") as f:
                f.write(kit_file)
            launch_args += [kit_file_path]
            carb.log_info(f"dump '{kit_file_path}':\n{kit_file}")

        launch_args += args

        launch_args += ["--exec", script_file + script_extra_args]

        await self._run_kit_and_check(launch_args, portable, timeout)

        if add_exec_check:
            self.assertTrue(os.path.exists(exec_check_file), "Script was not fully executed.")

    async def setUp(self):
        self._temp_folder = tempfile.mkdtemp()

    async def tearDown(self):
        shutil.rmtree(self._temp_folder)
