import os
import re
import sys

import click

from .sql_tester import SqlTester, TestCase, WorkPath


@click.group()
def main():
    pass


def create_backend(backend: str, env: str, case: TestCase, work_path: WorkPath):
    if backend == "spark":
        from easy_sql.local_spark import LocalSpark
        from easy_sql.sql_processor.backend import SparkBackend

        custom_conf = {
            "spark.sql.warehouse.dir": f"/tmp/spark-warehouse-{env}",
            "spark.driver.extraJavaOptions": (
                f"-Dderby.system.home=/tmp/spark-metastore-{env} "
                f"-Dderby.stream.error.file=/tmp/spark-metastore-{env}.log"
            ),
        }
        if case.udf_file_paths or case.func_file_paths:
            paths = case.udf_file_paths + case.func_file_paths
            custom_conf["spark.files"] = ",".join([work_path.path(path) for path in paths])
            # need to recreate SparkSession when there's new config
            LocalSpark.stop()

        import shutil

        shutil.rmtree(f"/tmp/spark-warehouse-{env}", True)

        # try get a clean spark environment
        spark = LocalSpark.get(custom_conf)
        backend = SparkBackend(spark)
        backend.clean()
        return backend
    elif backend == "postgres":
        from easy_sql.sql_processor.backend.rdb import RdbBackend

        assert "PG_URL" in os.environ, "Must set PG_URL env var to run an ETL with postgres backend."
        pg_url = os.environ["PG_URL"]
        backend = RdbBackend(pg_url)

        i = 0
        while True:
            try:
                backend.exec_native_sql(f"drop database if exists unittest{i}")
                break
            except Exception as e:
                print(e.args)
                if re.match(r".*psycopg2.errors.ObjectInUse.*", e.args[0]):
                    i += 1
                else:
                    raise e

        backend.exec_native_sql(f"create database unittest{i}")
        backend = RdbBackend(pg_url[: pg_url.rindex("/")] + f"/unittest{i}")
        return backend
    elif backend == "clickhouse":
        from easy_sql.sql_processor.backend.rdb import RdbBackend

        assert "CLICKHOUSE_URL" in os.environ, "Must set CLICKHOUSE_URL env var to run an ETL with clickhouse backend."
        ch_url = os.environ["CLICKHOUSE_URL"]
        backend = RdbBackend(ch_url)
        return backend
    elif backend == "bigquery":
        from easy_sql.sql_processor.backend.rdb import RdbBackend

        backend = RdbBackend(
            "bigquery://", credentials=f"{os.environ.get('HOME', '/tmp')}/.bigquery/credential-test.json"
        )
        return backend
    else:
        raise Exception("unsupported backend: " + backend)


_FILES_ARG_HELP = "Test files (multiple files split by comma): file1,file2,..."
_BACKEND_ARG_HELP = (
    "Which backend to use. The supported backends are: spark postgres bigquery clickhouse.\n"
    "For spark backend, the temporary folder will be /tmp/spark-metastore-{ENV} .\n"
    "For postgres backend, an environment variable named PG_URL must be set.\n"
    "For clickhouse backend, an environment variable named CLICKHOUSE_URL must be set.\n"
    "For bigquery backend, the credential file must be put at ${HOME}/.bigquery/credential-{ENV}.json .\n"
)
_ENV_ARG_HELP = (
    "Environment identifier to indicate which env to run on. Possible values could be: test e2e\n"
    "Used for logging and temporary file or folder creation."
)


@main.command(name="run-test")
@click.option("--files", "-f", type=str, help=_FILES_ARG_HELP)
@click.option("--backend", "-b", type=str, help=_BACKEND_ARG_HELP, required=False, default="spark", show_default=True)
@click.option("--env", "-e", type=str, help=_ENV_ARG_HELP, required=False, default="test", show_default=True)
def run_test(files, backend: str = "spark", env: str = "test"):
    _run_test(files, backend, env)


def _run_test(files, backend: str = "spark", env: str = "test"):
    test_data_files = [os.path.abspath(f.strip()) for f in files.split(",") if f.strip()]
    for f in test_data_files:
        if not os.path.exists(f):
            print(f"[ERROR] File does not exist: `{sys.argv}`")
            sys.exit(1)

    work_dir = os.path.abspath(os.curdir)
    SqlTester(
        env=env, backend_creator=lambda case: create_backend(backend, env, case, WorkPath(work_dir)), work_dir=work_dir
    ).run_tests(test_data_files)


@main.command(name="convert-json")
@click.option("--files", "-f", type=str, help=_FILES_ARG_HELP)
def convert_json(files):
    _convert_json(files)


def _convert_json(files):
    test_data_files = [os.path.abspath(f.strip()) for f in files.split(",") if f.strip()]
    sql_tester = SqlTester(work_dir=os.path.abspath(os.curdir))
    for f in test_data_files:
        sql_tester.convert_cases_to_json(f)


if __name__ == "__main__":
    main()
