from typing import Any, Union, Dict
from typing import Generic, TypeVar, Optional

import os
import json
from datetime import timedelta
import logging
import atexit
import time

from time import sleep
from windmill_api.models.whoami_response_200 import WhoamiResponse200

from windmill_api.client import AuthenticatedClient

from enum import Enum

from windmill_api.types import Unset

S = TypeVar("S")


class Resource(Generic[S]):
    pass


class JobStatus(Enum):
    WAITING = 1
    RUNNING = 2
    COMPLETED = 3


_client: "AuthenticatedClient | None" = None

logger = logging.getLogger("wmill_client")


def create_client(base_url: "str | None" = None, token: "str | None" = None) -> AuthenticatedClient:
    env_base_url = os.environ.get("BASE_INTERNAL_URL")

    if env_base_url is not None:
        env_base_url = env_base_url + "/api"

    base_url_: str = base_url or env_base_url or "http://localhost:8000/api"
    token_: str = token or os.environ.get("WM_TOKEN") or ""
    global _client
    if _client is None:
        _client = AuthenticatedClient(base_url=base_url_, token=token_, timeout=30, verify_ssl=False)
    return _client


def get_workspace() -> str:
    from_env = os.environ.get("WM_WORKSPACE")
    if from_env is None:
        raise Exception("Workspace not passed as WM_WORKSPACE")
    return from_env


def get_version() -> str:
    """
    Returns the current version of the backend
    """
    from windmill_api.api.settings import backend_version

    return backend_version.sync_detailed(client=create_client()).content.decode("us-ascii")


def run_script_async(
    hash: str,
    args: Dict[str, Any] = {},
    scheduled_in_secs: Union[None, int] = None,
) -> str:
    """
    Launch the run of a script and return immediately its job id
    """
    from windmill_api.api.job import run_script_by_hash

    from windmill_api.models.run_script_by_hash_json_body import RunScriptByHashJsonBody

    return run_script_by_hash.sync_detailed(
        client=create_client(),
        workspace=get_workspace(),
        hash_=hash,
        json_body=RunScriptByHashJsonBody.from_dict(args),
        scheduled_in_secs=scheduled_in_secs,
        parent_job=os.environ.get("DT_JOB_ID"),
    ).content.decode("us-ascii")


def run_script_sync(
    hash: str, 
    args: Optional[Dict[str, Any]] = None, 
    verbose: bool = False,
    assert_result_is_not_none: bool = True,
    cleanup: bool = True,
    timeout: Optional[timedelta] = None,
) -> Dict[str, Any]:
    """
    Run a script, wait for it to complete and return the result of the launched script
    """
    args = args or {}
    job_id = run_script_async(hash, args, None)

    def cancel_job():
        from windmill_api.api.job.cancel_queued_job import (
            sync_detailed,
            CancelQueuedJobJsonBody,
        )
        logger.warning(f"cancelling job {job_id}")
        return sync_detailed(
            workspace=get_workspace(),
            id=job_id,
            client=create_client(),
            json_body=CancelQueuedJobJsonBody(reason="killed by exit handler"),
        )

    if cleanup:
        atexit.register(cancel_job)

    nb_iter = 0

    start_time = time.time()
    timeout_seconds = timeout.total_seconds() if timeout else None

    while get_job_status(job_id) != JobStatus.COMPLETED:
        if timeout_seconds is not None:
            elapsed_time = time.time() - start_time
            if elapsed_time > timeout_seconds:
                msg = f"Script execution timed out after {timeout_seconds} seconds"
                logger.warning(msg)
                raise TimeoutError(msg)
        if verbose:
            print(f"Waiting for {job_id} to complete...")
        if nb_iter < 10:
            sleep(2.0)
        else:
            sleep(5.0)
        nb_iter += 1
    
    result = get_result(
        job_id,
        assert_result_is_not_none=assert_result_is_not_none,
    )

    # the job finished--we don't need to cancel it anymore
    if cleanup:
        atexit.unregister(cancel_job)

    error = isinstance(result, dict) and result.get("error")

    assert not error, error

    return result


def run_script_by_path_async(
    path: str,
    args: Dict[str, Any] = {},
    scheduled_in_secs: Union[None, int] = None,
) -> str:
    """
    Launch the run of a script and return immediately its job id
    """
    from windmill_api.api.job import run_script_by_path

    from windmill_api.models.run_script_by_path_json_body import RunScriptByPathJsonBody

    return run_script_by_path.sync_detailed(
        client=create_client(),
        workspace=get_workspace(),
        path=path,
        json_body=RunScriptByPathJsonBody.from_dict(args),
        scheduled_in_secs=scheduled_in_secs,
        parent_job=os.environ.get("DT_JOB_ID"),
    ).content.decode("us-ascii")


def run_script_by_path_sync(
    path: str, 
    args: Dict[str, Any] = {}, 
    verbose: bool = False,
    assert_result_is_not_none: bool = True,
    cleanup: bool = True,
    timeout: Optional[timedelta] = None,
) -> Dict[str, Any]:
    """
    Run a script, wait for it to complete and return the result of the launched script
    """
    args = args or {}
    job_id = run_script_by_path_async(path, args, None)

    def cancel_job():
        from windmill_api.api.job.cancel_queued_job import (
            sync_detailed,
            CancelQueuedJobJsonBody,
        )
        logger.warning(f"cancelling job {job_id}")
        return sync_detailed(
            workspace=get_workspace(),
            id=job_id,
            client=create_client(),
            json_body=CancelQueuedJobJsonBody(reason="killed by exit handler"),
        )

    if cleanup:
        atexit.register(cancel_job)

    nb_iter = 0

    start_time = time.time()
    timeout_seconds = timeout.total_seconds() if timeout else None

    while get_job_status(job_id) != JobStatus.COMPLETED:
        if timeout_seconds is not None:
            elapsed_time = time.time() - start_time
            if elapsed_time > timeout_seconds:
                msg = f"Script execution timed out after {timeout_seconds} seconds"
                logger.warning(msg)
                raise TimeoutError(msg)
        if verbose:
            print(f"Waiting for {job_id} to complete...")
        if nb_iter < 10:
            sleep(2.0)
        else:
            sleep(5.0)
        nb_iter += 1
    
    result = get_result(
        job_id,
        assert_result_is_not_none=assert_result_is_not_none,
    )

    # the job finished--we don't need to cancel it anymore
    if cleanup:
        atexit.unregister(cancel_job)

    error = isinstance(result, dict) and result.get("error")

    assert not error, error

    return result


def get_job_status(job_id: str) -> JobStatus:
    """
    Returns the status of a queued or completed job
    """
    from windmill_api.models.get_job_response_200_type import GetJobResponse200Type
    from windmill_api.api.job import get_job

    res = get_job.sync_detailed(client=create_client(), workspace=get_workspace(), id=job_id).parsed
    if not res:
        raise Exception(f"Job {job_id} not found")
    elif not res.type:
        raise Exception(f"Unexpected type not found for job {job_id}")
    elif res.type == GetJobResponse200Type.COMPLETEDJOB:
        return JobStatus.COMPLETED
    else:
        if not "running" in res.additional_properties:
            raise Exception(f"Unexpected running not found for completed job {job_id}")
        elif bool(res.additional_properties["running"]):
            return JobStatus.RUNNING
        else:
            return JobStatus.WAITING


def get_result(job_id: str, assert_result_is_not_none: bool = True) -> Dict[str, Any]:
    """
    Returns the result of a completed job
    """
    from windmill_api.api.job import get_completed_job

    res = get_completed_job.sync_detailed(client=create_client(), workspace=get_workspace(), id=job_id).parsed
    if not res:
        raise Exception(f"Job {job_id} not found")
    if assert_result_is_not_none and res.result is None:
        raise Exception(f"result was null for completed job {job_id}")
    else:
        return res.result


def get_resource(path: Union[str, None] = None, none_if_undefined: bool = False) -> Any:
    """
    Returns the resource at a given path
    """
    from windmill_api.api.resource import (
        get_resource_value_interpolated as get_resource_api,
    )

    path = path or get_state_path()
    parsed = get_resource_api.sync_detailed(workspace=get_workspace(), path=path, client=create_client())
    try:
        content = parsed.content.decode("utf-8")
        parsed = json.loads(content)
    except:
        parsed = None

    if parsed is None:
        if none_if_undefined:
            return None
        else:
            raise Exception(
                f"Resource at path {path} does not exist or you do not have read permissions on it: {content}"
            )

    return parsed


def duckdb_connection_settings(s3_resource: Any, none_if_undefined: bool = False) -> Union[str, None]:
    """
    Convenient helpers that takes an S3 resource as input and returns the settings necessary to
    initiate an S3 connection from DuckDB
    """
    from windmill_api.api.helpers import duckdb_connection_settings
    from windmill_api.models.s3_resource import S3Resource

    parsed = duckdb_connection_settings.sync_detailed(
        workspace=get_workspace(),
        client=create_client(),
        json_body={
            "s3_resource": s3_resource,
        },
    )
    try:
        content = parsed.content.decode("utf-8")
        parsed = json.loads(content)
    except:
        parsed = None

    if parsed is None:
        if none_if_undefined:
            return None
        else:
            raise Exception(f"Could not generate DuckDB S3 connection settings from the provided resource")
    return parsed


def polars_connection_settings(s3_resource: Any, none_if_undefined: bool = False) -> Any:
    """
    Convenient helpers that takes an S3 resource as input and returns the settings necessary to
    initiate an S3 connection from Polars
    """
    from windmill_api.api.helpers import polars_connection_settings

    parsed = polars_connection_settings.sync_detailed(
        workspace=get_workspace(),
        client=create_client(),
        json_body={
            "s3_resource": s3_resource,
        },
    )
    try:
        content = parsed.content.decode("utf-8")
        parsed = json.loads(content)
    except:
        parsed = None

    if parsed is None:
        if none_if_undefined:
            return None
        else:
            raise Exception(f"Could not generate Polars S3 connection settings from the provided resource")
    return parsed


def whoami() -> Union[WhoamiResponse200, None]:
    """
    Returns the current user
    """
    from windmill_api.api.user import whoami

    return whoami.sync(client=create_client(), workspace=get_workspace())


def get_state() -> Any:
    """
    Get the state
    """
    return get_resource(None, True)


def set_resource(value: Any, path: Union[str, None] = None, resource_type: str = "state") -> None:
    """
    Set the resource at a given path as a string, creating it if it does not exist
    """
    from windmill_api.models.create_resource_json_body import CreateResourceJsonBody
    from windmill_api.models.update_resource_value_json_body import (
        UpdateResourceValueJsonBody,
    )
    from windmill_api.api.resource import (
        exists_resource,
        update_resource_value,
        create_resource,
    )

    path = path or get_state_path()
    workspace = get_workspace()
    client = create_client()
    if not exists_resource.sync_detailed(workspace=workspace, path=path, client=client).parsed:
        create_resource.sync_detailed(
            workspace=workspace,
            client=client,
            json_body=CreateResourceJsonBody(path=path, value=value, resource_type=resource_type),
        )
    else:
        update_resource_value.sync_detailed(
            workspace=get_workspace(),
            client=client,
            path=path,
            json_body=UpdateResourceValueJsonBody(value=value),
        )


def set_state(value: Any) -> None:
    """
    Set the state
    """
    set_resource(value, None)


def set_shared_state_pickle(value: Any, path: str = "state.pickle") -> None:
    """
    Set the state in the shared folder using pickle
    """
    import pickle

    with open(f"/shared/{path}", "wb") as handle:
        pickle.dump(value, handle, protocol=pickle.HIGHEST_PROTOCOL)


def get_shared_state_pickle(path: str = "state.pickle") -> Any:
    """
    Get the state in the shared folder using pickle
    """
    import pickle

    with open(f"/shared/{path}", "rb") as handle:
        return pickle.load(handle)


def set_shared_state(value: Any, path: str = "state.json") -> None:
    """
    Set the state in the shared folder using pickle
    """
    import json

    with open(f"/shared/{path}", "w", encoding="utf-8") as f:
        json.dump(value, f, ensure_ascii=False, indent=4)


def get_shared_state(path: str = "state.json") -> None:
    """
    Set the state in the shared folder using pickle
    """
    import json

    with open(f"/shared/{path}", "r", encoding="utf-8") as f:
        return json.load(f)


def get_variable(path: str) -> str:
    """
    Returns the variable at a given path as a string
    """
    from windmill_api.api.variable import get_variable_value as get_variable_api

    res = get_variable_api.sync_detailed(workspace=get_workspace(), path=path, client=create_client())
    parsed = res.parsed
    if parsed is None:
        raise Exception(
            f"Variable at path {path} does not exist or you do not have read permissions on it: {res.content.decode('utf-8')}"
        )

    return parsed


def set_variable(path: str, value: str) -> None:
    """
    Set the variable at a given path as a string, creating it if it does not exist
    """
    from windmill_api.api.variable import (
        exists_variable,
        update_variable,
        create_variable,
    )

    from windmill_api.models.update_variable_json_body import UpdateVariableJsonBody
    from windmill_api.models.create_variable_json_body import CreateVariableJsonBody

    workspace = get_workspace()
    client = create_client()
    if not exists_variable.sync_detailed(workspace=workspace, path=path, client=client).parsed:
        create_variable.sync_detailed(
            workspace=workspace,
            client=client,
            json_body=CreateVariableJsonBody(path=path, value=value, is_secret=False, description=""),
        )
    else:
        update_variable.sync_detailed(
            workspace=get_workspace(),
            path=path,
            client=client,
            json_body=UpdateVariableJsonBody(value=value),
        )


def get_state_path() -> str:
    state_path = os.environ.get("WM_STATE_PATH_NEW") or os.environ.get("WM_STATE_PATH")
    if state_path is None:
        raise Exception("State path not found")
    return state_path


def get_resume_urls(approver: Union[str, None] = None) -> Dict:
    from windmill_api.api.job import get_resume_urls as get_resume_urls_api

    workspace = get_workspace()
    client = create_client()
    job_id = os.environ.get("WM_JOB_ID") or "NO_ID"
    import random

    nonce = random.randint(0, 1000000000)
    res = get_resume_urls_api.sync_detailed(workspace, job_id, nonce, client=client, approver=approver)
    if res.parsed is not None:
        return res.parsed.to_dict()
    else:
        raise Exception("Failed to get resume urls")
