'''
RunPod | CLI | Project | Functions
'''

import os
import sys
import uuid
from datetime import datetime

import tomlkit
from tomlkit import document, comment, table, nl

from runpod import __version__, get_pod, create_template, create_endpoint, update_endpoint_template
from runpod.cli import BASE_DOCKER_IMAGE, GPU_TYPES, ENV_VARS
from runpod.cli.utils.ssh_cmd import SSHConnection
from .helpers import (
    get_project_pod, copy_template_files,
    attempt_pod_launch, load_project_config, get_project_endpoint
)
from ...utils.rp_sync import sync_directory

STARTER_TEMPLATES = os.path.join(os.path.dirname(__file__), 'starter_templates')


def _launch_dev_pod():
    """ Launch a development pod. """
    config = load_project_config()  # Load runpod.toml

    print("Deploying development pod on RunPod...")

    # Prepare the environment variables
    environment_variables = {"RUNPOD_PROJECT_ID": config["project"]["uuid"]}
    for variable in config['project'].get('env_vars', {}):
        environment_variables[variable] = config['project']['env_vars'][variable]

    # Prepare the GPU types
    selected_gpu_types = config['project'].get('gpu_types', [])
    if config['project'].get('gpu', None):
        selected_gpu_types.append(config['project']['gpu'])

    # Attempt to launch a pod with the given configuration
    new_pod = attempt_pod_launch(config, environment_variables)
    if new_pod is None:
        print("Selected GPU types unavailable, try again later or use a different type.")
        return None

    print("Waiting for pod to come online... ", end="")
    sys.stdout.flush()

    # Wait for the pod to come online
    while new_pod.get('desiredStatus', None) != 'RUNNING' or new_pod.get('runtime') is None:
        new_pod = get_pod(new_pod['id'])

    project_pod_id = new_pod['id']

    print(f"Project {config['project']['name']} pod ({project_pod_id}) created.", end="\n\n")
    return project_pod_id


# -------------------------------- New Project ------------------------------- #
def create_new_project(project_name, runpod_volume_id, cuda_version, python_version,  # pylint: disable=too-many-locals, too-many-arguments, too-many-statements
                       model_type=None, model_name=None, init_current_dir=False):
    """ Create a new project. """
    if not init_current_dir:
        project_folder = os.path.join(os.getcwd(), project_name)
        if not os.path.exists(project_folder):
            os.makedirs(project_folder)

        if model_type is None:
            model_type = "default"

        template_dir = os.path.join(STARTER_TEMPLATES, model_type)

        copy_template_files(template_dir, project_folder)

        # Replace placeholders in requirements.txt
        requirements_path = os.path.join(project_folder, "builder/requirements.txt")
        with open(requirements_path, 'r', encoding='utf-8') as requirements_file:
            requirements_content = requirements_file.read()

        if "dev" in __version__:
            requirements_content = requirements_content.replace(
                '<<RUNPOD>>', 'git+https://github.com/runpod/runpod-python.git')
        else:
            requirements_content = requirements_content.replace(
                '<<RUNPOD>>', f'runpod=={__version__}')

        with open(requirements_path, 'w', encoding='utf-8') as requirements_file:
            requirements_file.write(requirements_content)

        # If there's a model_name, replace placeholders in handler.py
        if model_name:
            handler_path = os.path.join(project_name, "src/handler.py")
            with open(handler_path, 'r', encoding='utf-8') as file:
                handler_content = file.read()
                handler_content = handler_content.replace('<<MODEL_NAME>>', model_name)

            with open(handler_path, 'w', encoding='utf-8') as file:
                file.write(handler_content)
    else:
        project_folder = os.getcwd()

    project_uuid = str(uuid.uuid4())[:8]

    toml_config = document()
    toml_config.add(comment('RunPod Project Configuration'))
    toml_config.add(nl())
    toml_config.add("title", project_name)

    project_table = table()
    project_table.add("uuid", project_uuid)
    project_table.add("name", project_name)
    project_table.add("base_image", BASE_DOCKER_IMAGE.format(cuda_version=cuda_version))
    project_table.add("gpu_types", GPU_TYPES)
    project_table.add("gpu_count", 1)
    project_table.add("storage_id", runpod_volume_id)
    project_table.add("volume_mount_path", "/runpod-volume")
    project_table.add("ports", "8080/http, 22/tcp")
    project_table.add("container_disk_size_gb", 10)
    project_table.add("env_vars", ENV_VARS)
    toml_config.add("project", project_table)

    template_table = table()
    template_table.add("model_type", str(model_type))
    template_table.add("model_name", str(model_name))
    toml_config.add("template", template_table)

    runtime_table = table()
    runtime_table.add("python_version", python_version)
    runtime_table.add("handler_path", "src/handler.py")
    runtime_table.add("requirements_path", "builder/requirements.txt")
    toml_config.add("runtime", runtime_table)

    with open(os.path.join(project_folder, "runpod.toml"), 'w', encoding="UTF-8") as config_file:
        tomlkit.dump(toml_config, config_file)


# ------------------------------- Start Project ------------------------------ #
def start_project():  # pylint: disable=too-many-locals, too-many-branches
    '''
    Start the project development environment from runpod.toml
    - Check if the project pod already exists.

    - If the project pod does not exist:
        # SSH into the pod and create a project folder within the volume
        # Create a folder in the volume for the project that matches the project uuid
        # Create a folder in the project folder that matches the project name
        # Copy the project files into the project folder
        # crate a virtual environment using the python version specified in the project config
        # install the requirements.txt file

    - If the project pod does exist:

    '''
    config = load_project_config()  # Load runpod.toml

    for config_item in config['project']:
        print(f'    - {config_item}: {config["project"][config_item]}')
    print("")

    project_pod_id = get_project_pod(config['project']['uuid'])

    # Check if the project pod already exists, if not create it.
    if not project_pod_id:
        project_pod_id = _launch_dev_pod()

    if project_pod_id is None:
        return

    with SSHConnection(project_pod_id) as ssh_conn:

        project_path_uuid = f'{config["project"]["volume_mount_path"]}/{config["project"]["uuid"]}'
        project_path_uuid_dev = os.path.join(project_path_uuid, 'dev')
        project_path_uuid_prod = os.path.join(project_path_uuid, 'prod')
        remote_project_path = os.path.join(project_path_uuid_dev, config["project"]["name"])

        # Create the project folder on the pod
        print(f'Checking pod project folder: {remote_project_path} on pod {project_pod_id}')
        ssh_conn.run_commands([f'mkdir -p {remote_project_path} {project_path_uuid_prod}'])

        # Copy local files to the pod project folder
        print(f'Syncing files to pod {project_pod_id}')
        ssh_conn.rsync(os.getcwd(), project_path_uuid_dev)

        # Create the virtual environment
        venv_path = os.path.join(project_path_uuid_dev, "venv")
        print(f'Activating Python virtual environment: {venv_path} on pod {project_pod_id}')
        commands = [
            f'python{config["runtime"]["python_version"]} -m virtualenv {venv_path}',
            f'source {venv_path}/bin/activate &&'
            f'cd {remote_project_path} &&'
            'python -m pip install --upgrade pip &&'
            f'python -m pip install -v --requirement {config["runtime"]["requirements_path"]}'
        ]
        ssh_conn.run_commands(commands)

        # Start the watcher and then start the API development server
        sync_directory(ssh_conn, os.getcwd(), project_path_uuid_dev)

        project_name = config["project"]["name"]
        pip_req_path = os.path.join(remote_project_path, config['runtime']['requirements_path'])
        handler_path = os.path.join(remote_project_path, config['runtime']['handler_path'])

        launch_api_server = [f'''
            pkill inotify

            function force_kill {{
                kill $1 2>/dev/null
                sleep 1

                if ps -p $1 > /dev/null; then
                    echo "Graceful kill failed, attempting SIGKILL..."
                    kill -9 $1 2>/dev/null
                    sleep 1

                    if ps -p $1 > /dev/null; then
                        echo "Failed to kill process with PID: $1"
                        exit 1
                    else
                        echo "Killed process with PID: $1 using SIGKILL"
                    fi

                else
                    echo "Killed process with PID: $1"
                fi
            }}

            function cleanup {{
                echo "Cleaning up..."
                force_kill $last_pid
            }}

            trap cleanup EXIT SIGINT

            if source {project_path_uuid_dev}/venv/bin/activate; then
                echo -e "- Activated virtual environment."
            else
                echo "Failed to activate virtual environment."
                exit 1
            fi

            if cd {project_path_uuid_dev}/{project_name}; then
                echo -e "- Changed to project directory."
            else
                echo "Failed to change directory."
                exit 1
            fi

            exclude_pattern='(__pycache__|\\.pyc$)'
            function update_exclude_pattern {{
                exclude_pattern='(__pycache__|\\.pyc$)'
                if [[ -f .runpodignore ]]; then
                    while IFS= read -r line; do
                        line=$(echo "$line" | tr -d '[:space:]')
                        [[ "$line" =~ ^#.*$ || -z "$line" ]] && continue # Skip comments and empty lines
                        exclude_pattern="${{exclude_pattern}}|(${{line}})"
                    done < .runpodignore
                    echo -e "- Ignoring files matching pattern: $exclude_pattern"
                fi
            }}
            update_exclude_pattern

            # Start the API server in the background, and save the PID
            python {handler_path} --rp_serve_api --rp_api_host="0.0.0.0" --rp_api_port=8080 --rp_api_concurrency=1 &
            last_pid=$!

            echo -e "- Started API server with PID: $last_pid" && echo ""
            echo "Connect to the API server at:"
            echo ">  https://$RUNPOD_POD_ID-8080.proxy.runpod.net" && echo ""

            while true; do
                if changed_file=$(inotifywait -q -r -e modify,create,delete --exclude "$exclude_pattern" {remote_project_path} --format '%w%f'); then
                    echo "Detected changes in: $changed_file"
                else
                    echo "Failed to detect changes."
                    exit 1
                fi

                force_kill $last_pid

                if [[ $changed_file == *"requirements"* ]]; then
                    echo "Installing new requirements..."
                    python -m pip install --upgrade pip && python -m pip install -r {pip_req_path}
                fi

                if [[ $changed_file == *".runpodignore"* ]]; then
                    update_exclude_pattern
                fi

                python {handler_path} --rp_serve_api --rp_api_host="0.0.0.0" --rp_api_port=8080 --rp_api_concurrency=1 &
                last_pid=$!

                echo "Restarted API server with PID: $last_pid"
            done
        ''']

        print("")
        print("Starting project development endpoint...")
        ssh_conn.run_commands(launch_api_server)


# ------------------------------ Deploy Project ------------------------------ #
def create_project_endpoint():
    """ Create a project endpoint.
    - Move code in dev to prod folder
    - TODO: git commit the diff from current state to new state
    - Create a serverless template for the project
    - Create a new endpoint using the template
    """
    config = load_project_config()
    project_pod_id = get_project_pod(config['project']['uuid'])

    # Check if the project pod already exists, if not create it.
    if not project_pod_id:
        project_pod_id = _launch_dev_pod()

    if project_pod_id is None:
        return None

    with SSHConnection(project_pod_id) as ssh_conn:
        project_path_uuid = f'{config["project"]["volume_mount_path"]}/{config["project"]["uuid"]}'
        project_path_uuid_prod = os.path.join(project_path_uuid, 'prod')
        remote_project_path = os.path.join(project_path_uuid_prod, config["project"]["name"])

        # Copy local files to the pod project folder
        ssh_conn.run_commands([f'mkdir -p {remote_project_path}'])
        print(f'Syncing files to pod {project_pod_id} prod')
        ssh_conn.rsync(os.getcwd(), project_path_uuid_prod)

        # Create the virtual environment
        venv_path = os.path.join(project_path_uuid_prod, 'venv')
        print(f'Activating Python virtual environment: {venv_path} on pod {project_pod_id}')
        commands = [
            f'python{config["runtime"]["python_version"]} -m venv {venv_path}',
            f'source {venv_path}/bin/activate &&'
            f'cd {remote_project_path} &&'
            'python -m pip install --upgrade pip &&'
            f'python -m pip install -v --requirement {config["runtime"]["requirements_path"]}'
        ]
        ssh_conn.run_commands(commands)
        ssh_conn.close()

    environment_variables = {}
    for variable in config['project']['env_vars']:
        environment_variables[variable] = config['project']['env_vars'][variable]

    # Construct the docker start command
    activate_cmd = f'. /runpod-volume/{config["project"]["uuid"]}/prod/venv/bin/activate'
    python_cmd = f'python -u /runpod-volume/{config["project"]["uuid"]}/prod/{config["project"]["name"]}/{config["runtime"]["handler_path"]}'  # pylint: disable=line-too-long
    docker_start_cmd = 'bash -c "' + activate_cmd + ' && ' + python_cmd + '"'

    project_endpoint_template = create_template(
        name=f'{config["project"]["name"]}-endpoint | {config["project"]["uuid"]} | {datetime.now()}',  # pylint: disable=line-too-long
        image_name=config['project']['base_image'],
        container_disk_in_gb=config['project']['container_disk_size_gb'],
        docker_start_cmd=docker_start_cmd,
        env=environment_variables, is_serverless=True
    )

    deployed_endpoint = get_project_endpoint(config['project']['uuid'])
    if not deployed_endpoint:
        deployed_endpoint = create_endpoint(
            name=f'{config["project"]["name"]}-endpoint | {config["project"]["uuid"]}',
            template_id=project_endpoint_template['id'],
            network_volume_id=config['project']['storage_id'],
        )
    else:
        deployed_endpoint = update_endpoint_template(
            endpoint_id=deployed_endpoint['id'],
            template_id=project_endpoint_template['id'],
        )

    # does user want to tear down and recreate workers immediately?

    return deployed_endpoint['id']
