import asyncio
import inspect
import os
import secrets
import shutil
import sys
import tempfile
from enum import Enum
from pathlib import Path
from functools import wraps
from http import HTTPStatus
from importlib import import_module
from shutil import copytree
from tempfile import mkdtemp
from types import ModuleType
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

import requests
import yaml
from jina import Flow

from .config import DEFAULT_TIMEOUT, get_jcloud_config

if TYPE_CHECKING:
    from fastapi import FastAPI

APP_NAME = 'langchain'
BABYAGI_APP_NAME = 'babyagi'
PDF_QNA_APP_NAME = 'pdfqna'
PANDAS_AI_APP_NAME = 'pandasai'
AUTOGPT_APP_NAME = 'autogpt'

ServingGatewayConfigFile = 'servinggateway_config.yml'
APP_LOGS_URL = "[https://cloud.jina.ai/](https://cloud.jina.ai/user/flows?action=detail&id={app_id}&tab=logs)"
PRICING_URL = "****{cph}**** ([Read about pricing here](https://github.com/jina-ai/langchain-serve#-pricing))"


def syncify(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        return asyncio.run(f(*args, **kwargs))

    return wrapper


def hubble_exists(name: str, secret: Optional[str] = None) -> bool:
    return (
        requests.get(
            url='https://api.hubble.jina.ai/v2/executor/getMeta',
            params={'id': name, 'secret': secret},
        ).status_code
        == HTTPStatus.OK
    )


def _add_to_path(lcserve_app: bool = False):
    # add current directory to the beginning of the path to prioritize local imports
    sys.path.insert(0, os.getcwd())

    if lcserve_app:
        # get all directories in the apps folder and add them to the path
        for app in os.listdir(os.path.join(os.path.dirname(__file__), 'apps')):
            if os.path.isdir(os.path.join(os.path.dirname(__file__), 'apps', app)):
                sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'apps', app))


def _get_parent_dir(modname: str, filename: str) -> str:
    parts = modname.split('.')
    parent_dir = os.path.dirname(filename)
    for _ in range(len(parts) - 1):
        parent_dir = os.path.dirname(parent_dir)
    return parent_dir


def _load_module_from_str(module_str: str):
    try:
        module = import_module(module_str)
    except ModuleNotFoundError:
        print(f'Could not find module {module_str}')
        sys.exit(1)
    except AttributeError:
        print(f'Could not find appdir for module {module_str}')
        sys.exit(1)
    except Exception as e:
        print(f'Unknown error: {e}')
        sys.exit(1)
    return module


def _load_app_from_fastapi_app_str(
    fastapi_app_str: str,
) -> Tuple['FastAPI', ModuleType]:
    from .backend.playground.utils.helper import (
        ImportFromStringError,
        import_from_string,
    )

    try:
        fastapi_app, module = import_from_string(fastapi_app_str)
    except ImportFromStringError as e:
        print(f'Could not import app from {fastapi_app_str}: {e}')
        sys.exit(1)

    return fastapi_app, module


def _any_websocket_route_in_app(app: 'FastAPI') -> bool:
    from fastapi.routing import APIWebSocketRoute

    return any(isinstance(r, APIWebSocketRoute) for r in app.routes)


def _any_websocket_router_in_module(module: ModuleType) -> bool:
    # Go through the module and find all functions decorated by `serving` decorator
    for _, func in inspect.getmembers(module, inspect.isfunction):
        if hasattr(func, '__ws_serving__'):
            return True

    return False


def get_uri(id: str, tag: str):
    import requests
    from hubble import Auth

    r = requests.get(
        f"https://apihubble.jina.ai/v2/executor/getMeta?id={id}&tag={tag}",
        headers={"Authorization": f"token {Auth.get_auth_token()}"},
    )
    _json = r.json()
    if _json is None:
        print(f'Could not find image with id {id} and tag {tag}')
        return
    _image_name = _json['data']['name']
    _user_name = _json['meta']['owner']['name']
    return f'jinaai+docker://{_user_name}/{_image_name}:{tag}'


def get_module_dir(
    module_str: str = None,
    fastapi_app_str: str = None,
    app_dir: str = None,
    lcserve_app: bool = False,
) -> Tuple[str, bool]:
    _add_to_path(lcserve_app=lcserve_app)

    if module_str is not None:
        _module = _load_module_from_str(module_str)
        _is_websocket = _any_websocket_router_in_module(_module)
        _module_dir = _get_parent_dir(modname=module_str, filename=_module.__file__)
    elif fastapi_app_str is not None:
        fastapi_app, _module = _load_app_from_fastapi_app_str(fastapi_app_str)
        _is_websocket = _any_websocket_route_in_app(fastapi_app)
        _module_dir = _get_parent_dir(
            modname=fastapi_app_str, filename=_module.__file__
        )

    # if app_dir is not None, return it
    if app_dir is not None:
        return app_dir, _is_websocket

    if not _module.__file__.endswith('.py'):
        print(f'Unknown file type for module {module_str}')
        sys.exit(1)

    return _module_dir, _is_websocket


def _remove_langchain_serve(tmpdir: str) -> None:
    _requirements_txt = 'requirements.txt'
    _pyproject_toml = 'pyproject.toml'

    # Remove langchain-serve itself from the requirements list as a fixed version might break things
    if os.path.exists(os.path.join(tmpdir, _requirements_txt)):
        with open(os.path.join(tmpdir, _requirements_txt), 'r') as f:
            reqs = f.read().splitlines()

        reqs = [r for r in reqs if not r.startswith("langchain-serve")]
        with open(os.path.join(tmpdir, _requirements_txt), 'w') as f:
            f.write('\n'.join(reqs))

    if os.path.exists(os.path.join(tmpdir, _pyproject_toml)):
        import toml

        with open(os.path.join(tmpdir, _pyproject_toml), 'r') as f:
            pyproject = toml.load(f)

        if 'tool' in pyproject and 'poetry' in pyproject['tool']:
            poetry = pyproject['tool']['poetry']
            if 'dependencies' in poetry:
                poetry['dependencies'] = {
                    k: v
                    for k, v in poetry['dependencies'].items()
                    if k != 'langchain-serve'
                }

            if 'dev-dependencies' in poetry:
                poetry['dev-dependencies'] = {
                    k: v
                    for k, v in poetry['dev-dependencies'].items()
                    if k != 'langchain-serve'
                }

        with open(os.path.join(tmpdir, _pyproject_toml), 'w') as f:
            toml.dump(pyproject, f)


def _handle_dependencies(reqs: Tuple[str], tmpdir: str):
    # Create the requirements.txt if requirements are given
    _requirements_txt = 'requirements.txt'
    _pyproject_toml = 'pyproject.toml'

    _existing_requirements = []
    # Get existing requirements and add the new ones
    if os.path.exists(os.path.join(tmpdir, _requirements_txt)):
        with open(os.path.join(tmpdir, _requirements_txt), 'r') as f:
            _existing_requirements = tuple(f.read().splitlines())

    _new_requirements = []
    if reqs is not None:
        for _req in reqs:
            if os.path.isdir(_req):
                if os.path.exists(os.path.join(_req, _requirements_txt)):
                    with open(os.path.join(_req, _requirements_txt), 'r') as f:
                        _new_requirements = f.read().splitlines()

                elif os.path.exists(os.path.join(_req, _pyproject_toml)):
                    # copy pyproject.toml to tmpdir
                    shutil.copyfile(
                        os.path.join(_req, _pyproject_toml),
                        os.path.join(tmpdir, _pyproject_toml),
                    )
            elif os.path.isfile(_req):
                # if it's a file and name is requirements.txt, read it
                if os.path.basename(_req) == _requirements_txt:
                    with open(_req, 'r') as f:
                        _new_requirements = f.read().splitlines()
                elif os.path.basename(_req) == _pyproject_toml:
                    # copy pyproject.toml to tmpdir
                    shutil.copyfile(_req, os.path.join(tmpdir, _pyproject_toml))
            else:
                _new_requirements.append(_req)

        _final_requirements = set(_existing_requirements).union(set(_new_requirements))
        with open(os.path.join(tmpdir, _requirements_txt), 'w') as f:
            f.write('\n'.join(_final_requirements))

    _remove_langchain_serve(tmpdir)


def _handle_dockerfile(tmpdir: str, version: str):
    # if file `lcserve.Dockefile` exists, use it
    _lcserve_dockerfile = 'lcserve.Dockerfile'
    if os.path.exists(os.path.join(tmpdir, _lcserve_dockerfile)):
        shutil.copyfile(
            os.path.join(tmpdir, _lcserve_dockerfile),
            os.path.join(tmpdir, 'Dockerfile'),
        )

        # read the Dockerfile and replace the version
        with open(os.path.join(tmpdir, 'Dockerfile'), 'r') as f:
            dockerfile = f.read()

        dockerfile = dockerfile.replace(
            'jinawolf/serving-gateway:${version}',
            f'jinawolf/serving-gateway:{version}',
        )

        if 'ENTRYPOINT' not in dockerfile:
            dockerfile = (
                dockerfile
                + '\nENTRYPOINT [ "jina", "gateway", "--uses", "config.yml" ]'
            )

        with open(os.path.join(tmpdir, 'Dockerfile'), 'w') as f:
            f.write(dockerfile)

    else:
        # Create the Dockerfile
        with open(os.path.join(tmpdir, 'Dockerfile'), 'w') as f:
            dockerfile = [
                f'FROM jinawolf/serving-gateway:{version}',
                'COPY . /appdir/',
                'RUN if [ -e /appdir/requirements.txt ]; then pip install -r /appdir/requirements.txt; fi',
                'ENTRYPOINT [ "jina", "gateway", "--uses", "config.yml" ]',
            ]
            f.write('\n\n'.join(dockerfile))


def _handle_config_yaml(tmpdir: str, name: str):
    # Create the config.yml
    with open(os.path.join(tmpdir, 'config.yml'), 'w') as f:
        config_dict = {
            'jtype': 'ServingGateway',
            'py_modules': ['lcserve/backend/__init__.py'],
            'metas': {
                'name': name,
            },
        }
        f.write(yaml.safe_dump(config_dict, sort_keys=False))


def _push_to_hubble(
    tmpdir: str, name: str, tag: str, platform: str, verbose: bool, public: bool
) -> str:
    from hubble.executor.hubio import HubIO
    from hubble.executor.parsers import set_hub_push_parser

    from .backend.playground.utils.helper import EnvironmentVarCtxtManager

    secret = secrets.token_hex(8)
    args_list = [
        tmpdir,
        '--tag',
        tag,
        '--no-usage',
        '--no-cache',
    ]
    if verbose:
        args_list.remove('--no-usage')
        args_list.append('--verbose')
    if not public:
        args_list.append('--secret')
        args_list.append(secret)
        args_list.append('--private')

    args = set_hub_push_parser().parse_args(args_list)

    if platform:
        args.platform = platform

    if hubble_exists(name, secret):
        args.force_update = name

    push_envs = (
        {'JINA_HUBBLE_HIDE_EXECUTOR_PUSH_SUCCESS_MSG': 'true'} if not verbose else {}
    )
    with EnvironmentVarCtxtManager(push_envs):
        gateway_id = HubIO(args).push().get('id')
        return gateway_id + ':' + tag


def push_app_to_hubble(
    module_dir: str,
    image_name=None,
    tag: str = 'latest',
    requirements: Tuple[str] = None,
    version: str = 'latest',
    platform: str = None,
    verbose: Optional[bool] = False,
    public: Optional[bool] = False,
) -> str:
    from .backend.playground.utils.helper import get_random_name

    tmpdir = mkdtemp()

    # Copy appdir to tmpdir
    copytree(module_dir, tmpdir, dirs_exist_ok=True)
    # Copy lcserve to tmpdir
    copytree(
        os.path.dirname(__file__), os.path.join(tmpdir, 'lcserve'), dirs_exist_ok=True
    )

    if image_name is None:
        image_name = get_random_name()
    _handle_dependencies(requirements, tmpdir)
    _handle_dockerfile(tmpdir, version)
    _handle_config_yaml(tmpdir, image_name)
    return _push_to_hubble(tmpdir, image_name, tag, platform, verbose, public)


def get_gateway_config_yaml_path() -> str:
    return os.path.join(os.path.dirname(__file__), ServingGatewayConfigFile)


def get_gateway_uses(id: str) -> str:
    if id is not None:
        if id.startswith('jinahub+docker') or id.startswith('jinaai+docker'):
            return id
    return f'jinahub+docker://{id}'


def get_existing_name(app_id: str) -> str:
    from jcloud.flow import CloudFlow

    from .backend.playground.utils.helper import asyncio_run_property

    flow_obj = asyncio_run_property(CloudFlow(flow_id=app_id).status)
    if (
        'spec' in flow_obj
        and 'jcloud' in flow_obj['spec']
        and 'name' in flow_obj['spec']['jcloud']
    ):
        return flow_obj['spec']['jcloud']['name']


def get_global_jcloud_args(app_id: str = None, name: str = APP_NAME) -> Dict:
    if app_id is not None:
        _name = get_existing_name(app_id)
        if _name is not None:
            name = _name

    return {
        'jcloud': {
            'name': name,
            'labels': {
                'app': APP_NAME,
            },
            'monitor': {
                'traces': {
                    'enable': True,
                },
                'metrics': {
                    'enable': True,
                    'host': 'http://opentelemetry-collector.monitor.svc.cluster.local',
                    'port': 4317,
                },
            },
        }
    }


def get_uvicorn_args() -> Dict:
    return {
        'uvicorn_kwargs': {
            'ws_ping_interval': None,
            'ws_ping_timeout': None,
        }
    }


def get_with_args_for_jcloud(cors: bool = True, envs: Dict = {}) -> Dict:
    return {
        'with': {
            'cors': cors,
            'extra_search_paths': ['/workdir/lcserve'],
            'env': envs or {},
            **get_uvicorn_args(),
        }
    }


def get_flow_dict(
    module_str: str = None,
    fastapi_app_str: str = None,
    jcloud: bool = False,
    port: int = 8080,
    name: str = APP_NAME,
    timeout: int = DEFAULT_TIMEOUT,
    app_id: str = None,
    gateway_id: str = None,
    is_websocket: bool = False,
    jcloud_config_path: str = None,
    cors: bool = True,
    env: str = None,
    lcserve_app: bool = False,
) -> Dict:
    if jcloud:
        jcloud_config = get_jcloud_config(
            config_path=jcloud_config_path, timeout=timeout, is_websocket=is_websocket
        )

    _envs = {}
    if env is not None:
        # read env file and load to _envs dict
        from dotenv import dotenv_values

        _envs = dict(dotenv_values(env))

    uses = get_gateway_uses(id=gateway_id) if jcloud else get_gateway_config_yaml_path()
    flow_dict = {
        'jtype': 'Flow',
        **(get_with_args_for_jcloud(cors, _envs) if jcloud else {}),
        'gateway': {
            'uses': uses,
            'uses_with': {
                'modules': [module_str] if module_str else [],
                'fastapi_app_str': fastapi_app_str or '',
                'lcserve_app': lcserve_app,
            },
            'port': [port],
            'protocol': ['websocket'] if is_websocket else ['http'],
            'env': _envs if _envs else {},
            **get_uvicorn_args(),
            **(jcloud_config.to_dict() if jcloud else {}),
        },
        **(get_global_jcloud_args(app_id=app_id, name=name) if jcloud else {}),
    }
    if os.environ.get("LCSERVE_TEST", False):
        if 'with' not in flow_dict:
            flow_dict['with'] = {}

        flow_dict['with'].update(
            {
                'metrics': True,
                'metrics_exporter_host': 'http://localhost',
                'metrics_exporter_port': 4317,
            }
        )
    return flow_dict


def get_flow_yaml(
    module_str: str = None,
    fastapi_app_str: str = None,
    jcloud: bool = False,
    port: int = 8080,
    name: str = APP_NAME,
    is_websocket: bool = False,
    cors: bool = True,
    jcloud_config_path: str = None,
    env: str = None,
    lcserve_app: bool = False,
) -> str:
    return yaml.safe_dump(
        get_flow_dict(
            module_str=module_str,
            fastapi_app_str=fastapi_app_str,
            port=port,
            name=name,
            is_websocket=is_websocket,
            cors=cors,
            jcloud=jcloud,
            jcloud_config_path=jcloud_config_path,
            env=env,
            lcserve_app=lcserve_app,
        ),
        sort_keys=False,
    )


class ExportKind(str, Enum):
    KUBERNETES = 'kubernetes'
    DOCKER_COMPOSE = 'docker-compose'


def export_app(
    module_str: str,
    fastapi_app_str: str,
    app_dir: str,
    path: str,
    kind: ExportKind,
    image_name=None,
    tag: str = 'latest',
    requirements: Tuple[str] = None,
    version: str = 'latest',
    platform: str = None,
    verbose: Optional[bool] = False,
    public: Optional[bool] = False,
    name: str = APP_NAME,
    timeout: int = DEFAULT_TIMEOUT,
    env: str = None,
):
    module_dir, is_websocket = get_module_dir(
        module_str=module_str,
        fastapi_app_str=fastapi_app_str,
        app_dir=app_dir,
    )

    gateway_id = push_app_to_hubble(
        module_dir=module_dir,
        image_name=image_name,
        tag=tag,
        requirements=requirements,
        version=version,
        platform=platform,
        verbose=verbose,
        public=public,
    )

    flow_dict = get_flow_dict(
        module_str=module_str,
        fastapi_app_str=fastapi_app_str,
        jcloud=True,
        port=8080,
        name=name,
        timeout=timeout,
        app_id=None,
        gateway_id=gateway_id,
        is_websocket=is_websocket,
        jcloud_config_path=None,
        cors=True,
        env=env,
        lcserve_app=False,
    )

    f: Flow = Flow.load_config(flow_dict)

    if kind == ExportKind.KUBERNETES:
        f.to_kubernetes_yaml(path)
    elif kind == ExportKind.DOCKER_COMPOSE:
        _path = Path(path)
        if _path.is_file() and _path.suffix in ['.yml', '.yaml']:
            f.to_docker_compose_yaml(path)
        elif _path.is_dir():
            f.to_docker_compose_yaml(os.path.join(path, 'docker-compose.yml'))
        else:
            raise ValueError('path must be a file or a directory')


async def deploy_app_on_jcloud(
    flow_dict: Dict, app_id: str = None, verbose: bool = False
) -> Tuple[str, str]:
    from .backend.playground.utils.helper import EnvironmentVarCtxtManager

    os.environ['JCLOUD_LOGLEVEL'] = 'INFO' if verbose else 'ERROR'

    from jcloud.flow import CloudFlow

    with tempfile.TemporaryDirectory() as tmpdir:
        flow_path = os.path.join(tmpdir, 'flow.yml')
        with open(flow_path, 'w') as f:
            yaml.safe_dump(flow_dict, f, sort_keys=False)

        deploy_envs = {'JCLOUD_HIDE_SUCCESS_MSG': 'true'} if not verbose else {}
        with EnvironmentVarCtxtManager(deploy_envs):
            if app_id is None:  # appid is None means we are deploying a new app
                jcloud_flow = await CloudFlow(path=flow_path).__aenter__()
                app_id = jcloud_flow.flow_id

            else:  # appid is not None means we are updating an existing app
                jcloud_flow = CloudFlow(path=flow_path, flow_id=app_id)
                await jcloud_flow.update()

        for k, v in jcloud_flow.endpoints.items():
            if k.lower() == 'gateway (http)' or k.lower() == 'gateway (websocket)':
                return app_id, v

    return None, None


async def get_app_status_on_jcloud(app_id: str):
    from jcloud.flow import CloudFlow
    from rich import box
    from rich.align import Align
    from rich.console import Console
    from rich.markdown import Markdown
    from rich.table import Table

    _t = Table(
        'Attribute',
        'Value',
        show_header=False,
        box=box.ROUNDED,
        highlight=True,
        show_lines=True,
    )

    def _add_row(
        key,
        value,
        bold_key: bool = False,
        bold_value: bool = False,
        center_align: bool = True,
    ):
        return _t.add_row(
            Align(f'[bold]{key}' if bold_key else key, vertical='middle'),
            Align(f'[bold]{value}[/bold]' if bold_value else value, align='center')
            if center_align
            else value,
        )

    console = Console()
    with console.status(f'[bold]Getting app status for [green]{app_id}[/green]'):
        app_details = await CloudFlow(flow_id=app_id).status
        if app_details is None:
            return

        if 'status' not in app_details:
            return

        def _get_endpoint(app):
            endpoints = app.get('endpoints', {})
            return list(endpoints.values())[0] if endpoints else ''

        def _replace_wss_with_https(endpoint: str):
            return endpoint.replace('wss://', 'https://')

        status: Dict = app_details['status']
        total_cph = app_details.get('CPH', {}).get('total', 0)
        endpoint = _get_endpoint(status)

        _add_row('App ID', app_id, bold_key=True, bold_value=True)
        _add_row('Phase', status.get('phase', ''))
        _add_row('Endpoint', endpoint)
        _add_row(
            'App logs',
            Markdown(APP_LOGS_URL.format(app_id=app_id), justify='center'),
        )
        _add_row(
            'Base credits (per hour)',
            Markdown(PRICING_URL.format(cph=total_cph), justify='center'),
        )
        _add_row('Swagger UI', _replace_wss_with_https(f'{endpoint}/docs'))
        _add_row('OpenAPI JSON', _replace_wss_with_https(f'{endpoint}/openapi.json'))
        console.print(_t)


async def list_apps_on_jcloud(phase: str, name: str):
    from jcloud.flow import CloudFlow
    from jcloud.helper import cleanup_dt, get_phase_from_response
    from rich import box, print
    from rich.console import Console
    from rich.table import Table

    _t = Table(
        'AppID',
        'Phase',
        'Endpoint',
        'Created',
        box=box.ROUNDED,
        highlight=True,
    )

    console = Console()
    with console.status(f'[bold]Listing all apps'):
        all_apps = await CloudFlow().list_all(
            phase=phase, name=name, labels=f'app={APP_NAME}'
        )
        if not all_apps:
            print('No apps found')
            return

        def _get_endpoint(app):
            endpoints = app.get('status', {}).get('endpoints', {})
            return list(endpoints.values())[0] if endpoints else ''

        for app in all_apps['flows']:
            _t.add_row(
                app['id'],
                get_phase_from_response(app),
                _get_endpoint(app),
                cleanup_dt(app['ctime']),
            )
        console.print(_t)


async def remove_app_on_jcloud(app_id: str) -> None:
    from jcloud.flow import CloudFlow
    from rich import print

    await CloudFlow(flow_id=app_id).__aexit__()
    print(f'App [bold][green]{app_id}[/green][/bold] removed successfully!')


class ImportFromStringError(Exception):
    pass


def load_local_df(module: str):
    from importlib import import_module

    _add_to_path()

    module_str, _, attrs_str = module.partition(":")
    if not module_str or not attrs_str:
        message = (
            'Import string "{import_str}" must be in format "<module>:<attribute>".'
        )
        raise ImportFromStringError(message.format(import_str=module))

    try:
        module = import_module(module_str)
    except ImportError as exc:
        if exc.name != module_str:
            raise exc from None
        message = 'Could not import module "{module_str}".'
        raise ImportFromStringError(message.format(module_str=module_str))

    instance = module
    try:
        for attr_str in attrs_str.split("."):
            instance = getattr(instance, attr_str)
    except AttributeError:
        message = 'Could not import attribute "{attr_str}" from module "{module_str}".'
        raise ImportFromStringError(
            message.format(attr_str=attr_str, module_str=module_str)
        )

    return instance


def update_requirements(path: str, requirements: List[str]) -> List[str]:
    if os.path.exists(path):
        with open(path) as f:
            requirements.extend(f.read().splitlines())

    return requirements


def remove_prefix(text, prefix):
    return text[len(prefix) :] if text.startswith(prefix) else text
