import os
import time
from typing import Optional, List, Dict

from johnsnowlabs import settings
from johnsnowlabs.auto_install.jsl_home import get_install_suite_from_jsl_home
from johnsnowlabs.auto_install.softwares import Software
from johnsnowlabs.py_models.install_info import InstallSuite
from johnsnowlabs.utils.enums import JvmHardwareTarget


def authenticate_enviroment_HC(suite: InstallSuite):
    """Set Secret environ variables for Spark Context"""
    if suite.secrets.HC_LICENSE:
        os.environ["SPARK_NLP_LICENSE"] = suite.secrets.HC_LICENSE
    os.environ["AWS_ACCESS_KEY_ID"] = suite.secrets.AWS_ACCESS_KEY_ID
    os.environ["AWS_SECRET_ACCESS_KEY"] = suite.secrets.AWS_SECRET_ACCESS_KEY


def authenticate_enviroment_OCR(suite: InstallSuite):
    """Set Secret environ variables for Spark Context"""
    if suite.secrets.OCR_LICENSE:
        os.environ["SPARK_NLP_LICENSE"] = suite.secrets.OCR_LICENSE
    os.environ["AWS_ACCESS_KEY_ID"] = suite.secrets.AWS_ACCESS_KEY_ID
    os.environ["AWS_SECRET_ACCESS_KEY"] = suite.secrets.AWS_SECRET_ACCESS_KEY


def authenticate_enviroment_HC_and_OCR(suite: InstallSuite):
    """Set Secret environ variables for Spark Context"""
    authenticate_enviroment_HC(suite)
    authenticate_enviroment_OCR(suite)


def retry(fun, max_tries=10):
    for i in range(max_tries):
        try:
            time.sleep(0.3)
            fun()
            break
        except Exception:
            continue


def start(
    # -- JSL-Auth Flows --
    # Browser Auth
    browser_login: bool = False,
    # JWT Token Auth
    access_token: Optional[str] = None,
    # JSON file Auth
    json_license_path: Optional[str] = None,
    # AWS Auth
    aws_access_key: Optional[str] = None,
    aws_key_id: Optional[str] = None,
    # Manual License specification Auth
    enterprise_nlp_secret: Optional[str] = None,
    visual_secret: Optional[str] = None,
    hc_license: Optional[str] = None,
    ocr_license: Optional[str] = None,
    fin_license: Optional[str] = None,
    leg_license: Optional[str] = None,
    # License usage & Caching
    remote_license_number: int = 0,
    local_license_number: int = 0,
    store_in_jsl_home: bool = True,
    # -- Spark Session Configs --
    spark_conf: Optional[Dict[str, str]] = None,
    master_url: str = "local[*]",
    jar_paths: List[str] = None,
    # spark_nlp corresponds to enterprise NLP jar
    spark_nlp: bool = True,
    nlp: bool = True,
    visual: bool = False,
    hardware_target: str = JvmHardwareTarget.cpu.value,
    model_cache_folder: str = None,
    create_jsl_home_if_missing: bool = True,
) -> "pyspark.sql.SparkSession":
    from pyspark.sql import SparkSession

    already_launched = False
    if (
        "_instantiatedSession" in dir(SparkSession)
        and SparkSession._instantiatedSession is not None
    ):
        print("Spark Session already created, some configs may not take.")
        already_launched = True
        if settings.on_databricks:
            print(
                "Looks like you are on databricks which has a session pre-started. jsl.start() is not required. "
            )

    launched_products: List[str] = []
    hardware_target = JvmHardwareTarget.from_str(hardware_target)

    # Get all Local Jar Paths, downloads them if missing
    suite = get_install_suite_from_jsl_home(
        create_jsl_home_if_missing=create_jsl_home_if_missing,
        only_jars=True,
        jvm_hardware_target=hardware_target,
        force_browser=browser_login,
        browser_login=browser_login,
        access_token=access_token,
        local_license_number=local_license_number,
        remote_license_number=remote_license_number,
        secrets_file=json_license_path,
        hc_license=hc_license,
        hc_secret=enterprise_nlp_secret,
        ocr_secret=visual_secret,
        ocr_license=ocr_license,
        aws_access_key=aws_access_key,
        aws_key_id=aws_key_id,
        fin_license=fin_license,
        leg_license=leg_license,
        store_in_jsl_home=store_in_jsl_home,
    )

    # Collect all local Jar Paths we have access to for the SparkSession
    jars = []
    if (
        spark_nlp
        and Software.spark_nlp.check_installed(None)
        and suite.nlp.get_java_path()
    ):
        jars.append(suite.nlp.get_java_path())
        import sparknlp

        launched_products.append(
            f"{Software.spark_nlp.logo}{Software.spark_nlp.name}=={sparknlp.version()}"
        )

    if suite.secrets or not settings.license_required:
        if (
            suite.hc
            and nlp
            and Software.spark_hc.check_installed(None)
            and suite.hc.get_java_path()
        ):
            jars.append(suite.hc.get_java_path())
            if settings.license_required:
                authenticate_enviroment_HC(suite)
            import sparknlp_jsl

            launched_products.append(
                f"{Software.spark_hc.logo}{Software.spark_hc.name}=={sparknlp_jsl.version()}"
            )

        if (
            suite.ocr
            and visual
            and Software.spark_ocr.check_installed(None)
            and suite.ocr.get_java_path()
        ):
            jars.append(suite.ocr.get_java_path())
            if settings.license_required:
                authenticate_enviroment_OCR(suite)
            import sparkocr

            launched_products.append(
                f"{Software.spark_ocr.logo}{Software.spark_ocr.name}=={sparkocr.version()}"
            )
    import pyspark

    launched_products.append(
        f"running on {Software.spark.logo}{Software.pyspark.name}=={pyspark.version.__version__}"
    )

    builder = SparkSession.builder.appName(
        f'{settings.spark_session_name} with Jars for: {", ".join(launched_products)}'
    ).master(master_url)

    if jar_paths:
        # Add user specified Jars
        jars += jar_paths
    default_conf = {
        "spark.driver.memory": "16G",
        "spark.serializer": "org.apache.spark.serializer.KryoSerializer",
        "spark.kryoserializer.buffer.max": "2000M",
        "spark.driver.maxResultSize": "2000M",
        "spark.jars": ",".join(jars),
    }

    if suite.ocr and suite.ocr.get_java_path():
        # is_spark_version_env('32')
        default_conf["spark.sql.optimizer.expression.nestedPruning.enabled"] = "false"
        default_conf["spark.sql.optimizer.nestedSchemaPruning.enabled"] = "false"
        default_conf["spark.sql.legacy.allowUntypedScalaUDF"] = "true"
        default_conf["spark.sql.repl.eagerEval.enabled"] = "true"

    for k, v in default_conf.items():
        builder.config(str(k), str(v))

    if model_cache_folder:
        if not spark_conf:
            spark_conf = {}
        spark_conf["spark.jsl.settings.pretrained.cache_folder"] = model_cache_folder

    if spark_conf:
        for k, v in spark_conf.items():
            builder.config(str(k), str(v))
    spark = builder.getOrCreate()

    if suite.hc and exist_in_jvm(
        "com.johnsnowlabs.util.start.registerListenerAndStartRefresh"
    ):
        spark._jvm.com.johnsnowlabs.util.start.registerListenerAndStartRefresh()
    if suite.ocr and exist_in_jvm(
        "com.johnsnowlabs.util.OcrStart.registerListenerAndStartRefresh"
    ):
        retry(spark._jvm.com.johnsnowlabs.util.OcrStart.registerListenerAndStartRefresh)

    from colorama import Fore

    if not already_launched:
        print(
            f'👌 Launched {Fore.LIGHTGREEN_EX + hardware_target.value} optimized{Fore.RESET} session with with: {", ".join(launched_products)}'
        )

    return spark


def exist_in_jvm(java_class):
    # Check if a class path is present in JVM.
    # Only call this after Spark Session is started
    from pyspark.ml.util import _jvm
    from py4j.java_gateway import UserHelpAutoCompletion

    java_obj = _jvm()
    for name in java_class.split("."):
        # Bug in P4J. Even if ClassPath does not Exist, JVM response is proto.SUCCESS_PACKAGE
        # But it should give exception
        java_obj = getattr(java_obj, name)

    if UserHelpAutoCompletion.KEY in dir(java_obj):
        return False
    return True
