import shutil
from collections.abc import Generator

import pytest
import yaml
from click.testing import CliRunner
from kedro.framework.startup import ProjectMetadata

from kedro_databricks.cli.deploy import _get_arg_value
from kedro_databricks.constants import DEFAULT_TARGET
from kedro_databricks.plugin import commands
from kedro_databricks.utils import get_targets, read_databricks_config
from tests.utils import reset_init


def _load_resources(metadata, required_files):
    resource_dir = metadata.project_path / "resources"
    assert resource_dir.exists(), "Resource directory not created"
    assert resource_dir.is_dir(), "Resource directory is not a directory"
    files = [p.name for p in resource_dir.rglob("*")]
    files.sort()
    assert files == required_files, f"Resource files not created: {', '.join(files)}"
    resources = []
    for p in files:
        with open(resource_dir / p) as f:
            resource = yaml.safe_load(f)
            resources.append(resource)
    return zip(resources, files)


def _reset_bundle(metadata):
    """Reset the bundle to its initial state."""
    bundle_dir = metadata.project_path / "resources"
    if bundle_dir.exists():
        shutil.rmtree(bundle_dir)


def _validate_resources(metadata, required_files, task_validator):
    """Validate the resources generated by the bundle command."""
    resources = _load_resources(
        metadata=metadata,
        required_files=required_files,
    )
    for resource, file_name in resources:
        assert resource.get("resources") is not None

        jobs = resource["resources"]["jobs"]
        assert len(jobs) == 1

        job = jobs.get(file_name.split(".")[0])
        assert job is not None

        tasks = job.get("tasks")
        assert tasks is not None
        task_validator(tasks)


@pytest.fixture(scope="module")
def init_kedro_project(
    cli_runner, metadata, custom_provider
) -> Generator[tuple[ProjectMetadata, CliRunner]]:
    reset_init(metadata)
    _reset_bundle(metadata)
    init_cmd = ["databricks", "init", "--provider", custom_provider]
    result = cli_runner.invoke(commands, init_cmd, obj=metadata)
    assert result.exit_code == 0, (result.exit_code, result.stdout, result.exception)
    assert metadata.project_path.exists(), "Project path not created"
    assert metadata.project_path.is_dir(), "Project path is not a directory"
    assert metadata.project_path / "databricks.yml", "Databricks config not created"

    databricks_config = read_databricks_config(metadata.project_path)
    targets = get_targets(databricks_config)
    for target in targets:
        override_path = metadata.project_path / "conf" / target / "databricks.yml"
        assert (
            override_path.exists()
        ), f"Resource Overrides at {override_path} does not exist"
    yield metadata, cli_runner
    reset_init(metadata)
    _reset_bundle(metadata)


def test_databricks_bundle_fail(cli_runner, metadata):
    bundle_fail = ["databricks", "bundle", "--default-key", "_deault"]
    result = cli_runner.invoke(commands, bundle_fail, obj=metadata)
    assert result.exit_code == 1, (result.exit_code, result.stdout, result.exception)


def test_databricks_bundle_with_overrides(init_kedro_project):
    metadata, cli_runner = init_kedro_project

    command = ["databricks", "bundle", "--env", DEFAULT_TARGET]
    result = cli_runner.invoke(commands, command, obj=metadata)
    assert result.exit_code == 0, (result.exit_code, result.stdout, result.exception)

    def task_validator(tasks):
        assert len(tasks) == 8
        for i, task in enumerate(tasks):
            assert task.get("task_key") in (f"node{i}", f"ns_{i}_node_{i}_1")
            assert task.get("job_cluster_key") == "default"
            params = task.get("python_wheel_task").get("parameters")
            assert params is not None
            env = _get_arg_value(params, "--env")
            assert env is not None
            assert env == "${var.environment}"

    _validate_resources(
        metadata=metadata,
        required_files=[
            f"{metadata.package_name}.yml",
            f"{metadata.package_name}_ds.yml",
            f"{metadata.package_name}_namespaced_pipeline.yml",
        ],
        task_validator=task_validator,
    )


def test_databricks_bundle_with_conf(init_kedro_project):
    """Test the `bundle` command"""
    metadata, cli_runner = init_kedro_project

    command = [
        "databricks",
        "bundle",
        "--env",
        DEFAULT_TARGET,
        "--conf-source",
        "conf/sub_pipeline",
    ]
    result = cli_runner.invoke(commands, command, obj=metadata)
    assert result.exit_code == 0, (result.exit_code, result.stdout, result.exception)

    def task_validator(tasks):
        assert len(tasks) == 8
        for i, task in enumerate(tasks):
            assert task.get("task_key") in (f"node{i}", f"ns_{i}_node_{i}_1")
            assert task.get("job_cluster_key") == "default"
            params = task.get("python_wheel_task").get("parameters")
            assert params is not None
            env = _get_arg_value(params, "--env")
            assert env is not None
            assert env == "${var.environment}"

    _validate_resources(
        metadata=metadata,
        required_files=[
            f"{metadata.package_name}.yml",
            f"{metadata.package_name}_ds.yml",
            f"{metadata.package_name}_namespaced_pipeline.yml",
        ],
        task_validator=task_validator,
    )


def test_databricks_bundle_without_overrides(init_kedro_project):
    """Test the `bundle` command"""
    metadata, cli_runner = init_kedro_project

    command = ["databricks", "bundle", "--env", DEFAULT_TARGET]
    result = cli_runner.invoke(commands, command, obj=metadata)
    assert result.exit_code == 0, (result.exit_code, result.stdout, result.exception)

    def task_validator(tasks):
        assert len(tasks) == 8
        for i, task in enumerate(tasks):
            assert task.get("task_key") in (f"node{i}", f"ns_{i}_node_{i}_1")
            assert task.get("job_cluster_key") == "default"
            params = task.get("python_wheel_task").get("parameters")
            assert params is not None
            env = _get_arg_value(params, "--env")
            assert env is not None
            assert env == "${var.environment}"

    _validate_resources(
        metadata=metadata,
        required_files=[
            f"{metadata.package_name}.yml",
            f"{metadata.package_name}_ds.yml",
            f"{metadata.package_name}_namespaced_pipeline.yml",
        ],
        task_validator=task_validator,
    )


def test_databricks_bundle_with_params(init_kedro_project):
    """Test the `bundle` command"""
    metadata, cli_runner = init_kedro_project
    command = [
        "databricks",
        "bundle",
        "--env",
        DEFAULT_TARGET,
        "--params",
        "run_date={{job.parameters.run_date}},run_id={{job.parameters.run_id}}",
        "--overwrite",
    ]
    result = cli_runner.invoke(commands, command, obj=metadata)
    assert result.exit_code == 0, (result.exit_code, result.stdout, result.exception)

    def task_validator(tasks):
        assert len(tasks) == 8
        for i, task in enumerate(tasks):
            assert task.get("task_key") in (f"node{i}", f"ns_{i}_node_{i}_1")
            assert task.get("job_cluster_key") == "default"
            params = task.get("python_wheel_task").get("parameters")
            assert params is not None
            value = _get_arg_value(params, "--params")
            assert value is not None
            assert (
                value
                == "run_date={{job.parameters.run_date}},run_id={{job.parameters.run_id}}"
            )

    _validate_resources(
        metadata=metadata,
        required_files=[
            f"{metadata.package_name}.yml",
            f"{metadata.package_name}_ds.yml",
            f"{metadata.package_name}_namespaced_pipeline.yml",
        ],
        task_validator=task_validator,
    )


def test_databricks_bundle_with_pipeline(init_kedro_project):
    """Test the `bundle` command"""
    metadata, cli_runner = init_kedro_project
    if (metadata.project_path / "resources").exists():
        shutil.rmtree(metadata.project_path / "resources")
    command = [
        "databricks",
        "bundle",
        "--env",
        DEFAULT_TARGET,
        "--pipeline",
        "ds",
        "--overwrite",
    ]
    result = cli_runner.invoke(commands, command, obj=metadata)
    assert result.exit_code == 0, (result.exit_code, result.stdout, result.exception)

    def task_validator(tasks):
        assert len(tasks) == 8
        for i, task in enumerate(tasks):
            assert task.get("task_key") in (f"node{i}", f"ns_{i}_node_{i}_1")
            assert task.get("job_cluster_key") == "default"
            params = task.get("python_wheel_task").get("parameters")
            assert params is not None
            env = _get_arg_value(params, "--env")
            assert env is not None
            assert env == "${var.environment}"
            pipeline = _get_arg_value(params, "--pipeline")
            assert pipeline is None

    _validate_resources(
        metadata=metadata,
        required_files=[
            f"{metadata.package_name}_ds.yml",
        ],
        task_validator=task_validator,
    )


def test_databricks_bundle_with_no_nodes(init_kedro_project):
    """Test the `bundle` command"""
    metadata, cli_runner = init_kedro_project
    command = [
        "databricks",
        "bundle",
        "--env",
        DEFAULT_TARGET,
        "--pipeline",
        "non-existing-pipeline",
        "--overwrite",
    ]
    result = cli_runner.invoke(commands, command, obj=metadata)
    assert result.exit_code == 1, (result.exit_code, result.stdout, result.exception)
    assert isinstance(result.exception, KeyError)
    exception = bytes(str(result.exception), "utf-8").decode("unicode_escape")
    expected = "Pipeline 'non-existing-pipeline' not found."
    assert (
        expected in result.exception.args[0]
    ), f"Expected exception message: {expected} not found in {exception}"
