"""
Remote dt CLI client for Datatailr.

Usage:
  - Run `datatailr login` to authenticate.
  - Use `dt <command> <subcommand> [args...]` as usual - the calls will be executed remotely.
  - To get help, either append -h / --help to a CLI call, or call bare `dt`, `dt <command>` or `dt <command> <subcommand>`.

Notes:
  - Files specified as arguments will be read from local disk and sent as JSON when possible.
  - This client requires the 'requests' library.
"""

from __future__ import annotations

import base64
import json
import sys
import tarfile
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import requests

from datatailr._remote_client.parser import build_parser
from datatailr._remote_client.ssh_setup import (
    select_workstation_interactively,
    setup_ssh_config,
)
from datatailr._remote_client.blob_handlers import (
    BLOB_MAX_SIZE,
    handle_blob_cp,
    handle_blob_get,
    handle_blob_put,
)
from datatailr._remote_client.utils import (
    Config,
    build_headers,
    interactive_login,
    is_auth_failure,
    load_packaged_api_json,
    validate_remote_session,
    username_from_jwt,
)


def get_non_option_args(subcommand: Dict[str, Any]) -> List[Dict[str, Any]]:
    non_options = subcommand.get("non_options") or {}
    arg_names = non_options.get("arg_names") or []
    arg_types = non_options.get("arg_types") or []
    required_original = {
        str(name)
        for name in (non_options.get("required") or [])
        if isinstance(name, str)
    }

    args: List[Dict[str, Any]] = []
    for index, arg_name in enumerate(arg_names):
        if not isinstance(arg_name, str) or not arg_name:
            continue
        raw_type = str(arg_types[index] if index < len(arg_types) else "String").lower()
        request_payload_key = "data" if raw_type == "file" else arg_name
        args.append(
            {
                "request_payload_key": request_payload_key,
                "required": arg_name in required_original,
            }
        )
    return args


def get_subcommand_http_method(subcommand: Dict[str, Any]) -> str:
    non_options = subcommand.get("non_options") or {}
    return str(non_options.get("http_method", "get")).upper()


def load_api_option_metadata(
    subcommand: Dict[str, Any],
) -> Tuple[Dict[str, str], Dict[str, bool]]:
    short_to_long: Dict[str, str] = {}
    long_requires: Dict[str, bool] = {}

    for arg in get_non_option_args(subcommand):
        long_requires[str(arg["request_payload_key"])] = True

    for option in subcommand.get("options") or []:
        if not isinstance(option, list) or len(option) < 3:
            continue
        short = option[0] if len(option) > 0 and isinstance(option[0], str) else ""
        long = option[1] if len(option) > 1 and isinstance(option[1], str) else ""
        kind = option[2] if len(option) > 2 and isinstance(option[2], str) else ""
        if not long:
            continue
        long_requires[long] = kind == "required_argument"
        if short:
            short_to_long[short] = long

    return short_to_long, long_requires


def parse_options_from_args(
    argv_tail: List[str],
    subcommand: Dict[str, Any],
    short_to_long: Dict[str, str],
    long_requires: Dict[str, bool],
) -> Tuple[List[str], Dict[str, Any]]:
    valid_long = set(long_requires.keys())
    for option in subcommand.get("options") or []:
        if not isinstance(option, list) or len(option) < 2:
            continue
        long = option[1]
        if isinstance(long, str) and long:
            valid_long.add(long)
    positionals: List[str] = []
    options: Dict[str, Any] = {}

    i = 0
    while i < len(argv_tail):
        token = argv_tail[i]
        if token == "--":
            positionals.extend(argv_tail[i + 1 :])
            break

        if token.startswith("--") and len(token) > 2:
            name = token[2:]
            option_value: Optional[str] = None
            if "=" in name:
                name, option_value = name.split("=", 1)
            if name not in valid_long and name not in long_requires:
                positionals.append(token)
                i += 1
                continue
            requires_arg = long_requires.get(name, True)
            if not requires_arg:
                options[name] = True
                i += 1
                continue
            if option_value is not None:
                options[name] = option_value
                i += 1
                continue
            if i + 1 >= len(argv_tail):
                raise RuntimeError(f"Missing value for option '--{name}'")
            options[name] = argv_tail[i + 1]
            i += 2
            continue

        if token.startswith("-") and len(token) >= 2:
            shorts = token[1:]
            handled = False
            for idx, short in enumerate(shorts):
                long_name = short_to_long.get(short)
                if not long_name:
                    handled = False
                    break
                requires_arg = long_requires.get(long_name, True)
                if requires_arg:
                    rest = shorts[idx + 1 :]
                    if rest:
                        options[long_name] = rest
                        handled = True
                        break
                    if i + 1 >= len(argv_tail):
                        raise RuntimeError(f"Missing value for option '-{short}'")
                    options[long_name] = argv_tail[i + 1]
                    i += 1
                    handled = True
                    break
                options[long_name] = True
                handled = True
            if handled:
                i += 1
                continue

        positionals.append(token)
        i += 1

    return positionals, options


def encode_file_as_json(path_str: str) -> Any:
    p = Path(path_str)
    text = p.read_text(encoding="utf-8")
    try:
        return json.loads(text)
    except Exception as exc:
        raise RuntimeError(f"Failed to parse file '{path_str}' as JSON: {exc}")


def resolve_repo_path(path_to_repo: str, job_file: Path) -> Path:
    repo_path = Path(path_to_repo)
    if repo_path.exists():
        return repo_path
    fallback = job_file.parent / repo_path.name
    if fallback.exists():
        return fallback
    raise RuntimeError(
        f"Cannot resolve path_to_repo '{path_to_repo}'. "
        f"Expected it to exist or be relative to '{job_file.parent}'."
    )


def create_module_archive(repo_path: Path, module_path: str) -> Tuple[str, str, str]:
    module_dir = repo_path / module_path
    if not module_dir.exists() or not module_dir.is_dir():
        raise RuntimeError(
            f"Module path '{module_path}' not found under repo '{repo_path}'."
        )
    repo_dir_name = str(repo_path).lstrip("/")
    if not repo_dir_name:
        raise RuntimeError(
            f"Cannot infer repo directory name from path_to_repo '{repo_path}'."
        )

    with tempfile.TemporaryDirectory() as tmp_dir:
        archive_path = Path(tmp_dir) / "module.tar.gz"
        with tarfile.open(archive_path, "w:gz") as tar:
            tar.add(module_dir, arcname=module_path)
        archive_b64 = base64.b64encode(archive_path.read_bytes()).decode("ascii")
    return archive_b64, repo_dir_name, module_path


def try_to_prepare_repo_bundle(job_json: Any, job_file: Path) -> Dict[str, Any]:
    if not isinstance(job_json, dict):
        return {}
    image = job_json.get("image", {})
    path_to_repo = image.get("path_to_repo")
    path_to_module = image.get("path_to_module")
    if not path_to_repo or not path_to_module:
        return {}

    repo_path = resolve_repo_path(str(path_to_repo), job_file)
    archive_b64, repo_dir_name, module_path = create_module_archive(
        repo_path, str(path_to_module)
    )
    import socket

    return {
        "repo_archive_b64": archive_b64,
        "repo_dir_name": repo_dir_name,
        "repo_module_path": module_path,
        "local_hostname": socket.gethostname(),
    }


def is_local_file_arg(value: str) -> bool:
    if value.startswith("blob://"):
        return False
    p = Path(value)
    return p.exists() and p.is_file()


def get_subcommand_definition(
    api_json: Dict[str, Any], cmd: str, subcmd: str
) -> Tuple[str, Dict[str, Any]]:
    command = api_json.get(cmd)
    if not isinstance(command, dict):
        raise RuntimeError(f"Unknown command group '{cmd}'")
    all_subcommands = command.get("sub_commands")
    if not isinstance(all_subcommands, dict):
        raise RuntimeError(f"Command group '{cmd}' has no subcommands")
    subcommand = all_subcommands.get(subcmd)
    if not isinstance(subcommand, dict):
        raise RuntimeError(f"Unknown subcommand '{cmd} {subcmd}'")
    return get_subcommand_http_method(subcommand), subcommand


def build_request_from_subcommand(
    subcommand_definition: Dict[str, Any],
    argv_tail: List[str],
    options: Optional[Dict[str, Any]] = None,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    query: Dict[str, Any] = {}
    body: Dict[str, Any] = {}
    method = get_subcommand_http_method(subcommand_definition)
    if method == "GET":
        query.update(options or {})
    else:
        body.update(options or {})
    target = query if method == "GET" else body

    required_names: List[str] = []
    optional_names: List[str] = []
    seen_required: set[str] = set()
    seen_optional: set[str] = set()
    for arg in get_non_option_args(subcommand_definition):
        key = str(arg["request_payload_key"])
        if key in target:
            continue
        if arg["required"]:
            if key not in seen_required:
                required_names.append(key)
                seen_required.add(key)
        elif key not in seen_optional:
            optional_names.append(key)
            seen_optional.add(key)

    required_query_names = required_names if method == "GET" else []
    required_body_names = required_names if method != "GET" else []

    cron_preview = bool(options and options.get("cron-string"))
    # Native dt job run/save -c does not consume positionals (no job file on that path).
    min_expected = 0 if cron_preview else len(required_names)
    if len(argv_tail) < min_expected:
        raise RuntimeError(
            f"Wrong number of arguments.\nExpected at least {min_expected} positional args "
            f"(required query={required_query_names}, required body={required_body_names}), "
            f"got {len(argv_tail)}."
        )

    def assign_value(name: str, value: str) -> None:
        clean_val = value[7:] if value.startswith("file://") else value
        if name == "data" and is_local_file_arg(clean_val):
            job_file = Path(clean_val)
            job_json = encode_file_as_json(clean_val)
            extra_bundle = try_to_prepare_repo_bundle(job_json, job_file)
            body[name] = job_json
            body.update(extra_bundle)
            return
        target[name] = value

    cursor = 0
    to_require = [] if cron_preview else required_names
    to_optional = [] if cron_preview else optional_names
    for name in to_require:
        assign_value(name, argv_tail[cursor])
        cursor += 1

    for name in to_optional:
        if cursor >= len(argv_tail):
            break
        assign_value(name, argv_tail[cursor])
        cursor += 1

    return query, body


def send_help(
    cfg: Config, group: Optional[str], subcmd: Optional[str], timeout_s: int = 30
) -> int:
    if group and subcmd:
        url = f"{cfg.base_url}/api/{group}/{subcmd}/help"
    elif group:
        url = f"{cfg.base_url}/api/{group}/help"
    else:
        url = f"{cfg.base_url}/api/help"

    try:
        resp = requests.get(url, headers=build_headers(cfg), timeout=timeout_s)
    except requests.RequestException:
        print(
            f"Cannot reach the platform at {cfg.base_url}. Is the platform running?\n"
            "If the URL is wrong, run 'datatailr login' to update it.",
            file=sys.stderr,
        )
        return 2

    return handle_response(resp)


def send_request(
    cfg: Config, api_json: Dict[str, Any], full_argv: List[str], timeout_s: int = 60
) -> int:
    def handle_unknown_command(target_cmd: str) -> int:
        cmd_exists = isinstance(api_json.get(target_cmd), dict)
        if cmd_exists:
            return send_help(cfg, target_cmd, None, timeout_s=timeout_s)
        return send_help(cfg, None, None, timeout_s=timeout_s)

    if not full_argv:
        return send_help(cfg, None, None, timeout_s=timeout_s)
    if len(full_argv) == 1:
        return handle_unknown_command(full_argv[0])

    cmd, subcmd, *rest = full_argv
    if rest == ["-h"] or rest == ["--help"]:
        return send_help(cfg, cmd, subcmd, timeout_s=timeout_s)
    api_path = f"/api/{cmd}/{subcmd}"

    if cmd == "blob" and subcmd in ("cp", "put", "get"):
        try:
            method, subcommand_definition = get_subcommand_definition(
                api_json, cmd, subcmd
            )
        except RuntimeError:
            return handle_unknown_command(cmd)
        short_to_long, long_requires = load_api_option_metadata(subcommand_definition)
        positionals, options = parse_options_from_args(
            rest, subcommand_definition, short_to_long, long_requires
        )
        handler = {
            "cp": handle_blob_cp,
            "put": handle_blob_put,
            "get": handle_blob_get,
        }[subcmd]
        return handler(cfg, positionals, options, send_help, handle_response, timeout_s)

    try:
        method, subcommand_definition = get_subcommand_definition(api_json, cmd, subcmd)
    except RuntimeError:
        return handle_unknown_command(cmd)

    short_to_long, long_requires = load_api_option_metadata(subcommand_definition)
    positionals, options = parse_options_from_args(
        rest, subcommand_definition, short_to_long, long_requires
    )

    try:
        query, body = build_request_from_subcommand(
            subcommand_definition, positionals, options=options
        )
    except RuntimeError as exc:
        if str(exc).startswith("Wrong number of arguments"):
            return send_help(cfg, cmd, subcmd, timeout_s=timeout_s)
        raise

    url = f"{cfg.base_url}{api_path}"
    headers = build_headers(cfg)

    try:
        if method == "GET":
            resp = requests.get(url, headers=headers, params=query, timeout=timeout_s)
        else:
            resp = requests.request(
                method,
                url,
                headers=headers,
                params=query,
                json=(body if body else None),
                timeout=timeout_s,
            )
    except requests.RequestException:
        print(
            f"Cannot reach the platform at {cfg.base_url}. Is the platform running?\n"
            "If the URL is wrong, run 'datatailr login' to update it.",
            file=sys.stderr,
        )
        return 2

    if cmd == "blob" and len(resp.content) > BLOB_MAX_SIZE:
        print(
            f"Response too large ({len(resp.content)} bytes). "
            f"Maximum allowed size for remote blob operations is {BLOB_MAX_SIZE} bytes (100 MB).",
            file=sys.stderr,
        )
        return 2

    return handle_response(resp, pretty_print=options.get("pretty", False))


def _parse_api_response_or_raise(resp: requests.Response) -> Any:
    if is_auth_failure(resp):
        raise RuntimeError(
            "Not authenticated or session expired - please call 'datatailr login' to authenticate."
        )

    if resp.status_code >= 400:
        try:
            data = resp.json()
            err = data.get("error", "")
            if err:
                raise RuntimeError(str(err))
            if "message" in data:
                raise RuntimeError(str(data["message"]))
        except RuntimeError:
            raise
        except Exception:
            pass
        raise RuntimeError(resp.text or f"HTTP {resp.status_code}")

    content_type = resp.headers.get("Content-Type", "").lower()
    if "application/json" in content_type:
        data = resp.json()
        if isinstance(data, dict):
            err = data.get("error", "")
            if err:
                raise RuntimeError(str(err))
            return data.get("result")
        return data

    if not resp.text:
        return None
    try:
        return json.loads(resp.text)
    except Exception:
        return resp.text


def _api_call(
    cfg: Config,
    api_json: Dict[str, Any],
    cmd: str,
    subcmd: str,
    argv_tail: List[str],
    timeout_s: int = 60,
) -> Any:
    api_path = f"/api/{cmd}/{subcmd}"
    method, subcommand_definition = get_subcommand_definition(api_json, cmd, subcmd)

    short_to_long, long_requires = load_api_option_metadata(subcommand_definition)
    positionals, options = parse_options_from_args(
        argv_tail, subcommand_definition, short_to_long, long_requires
    )
    query, body = build_request_from_subcommand(
        subcommand_definition, positionals, options=options
    )

    url = f"{cfg.base_url}{api_path}"
    headers = build_headers(cfg)
    if method == "GET":
        resp = requests.get(url, headers=headers, params=query, timeout=timeout_s)
    else:
        resp = requests.request(
            method,
            url,
            headers=headers,
            params=query,
            json=(body if body else None),
            timeout=timeout_s,
        )

    return _parse_api_response_or_raise(resp)


def handle_response(resp: requests.Response, pretty_print: bool = False) -> int:
    if is_auth_failure(resp):
        print(
            "Not authenticated or session expired - please call 'datatailr login' to authenticate.",
            file=sys.stderr,
        )
        return 1

    if "license.html" in resp.url:
        print(
            "License for this Datatailr installation has expired. Please contact your administrator.",
            file=sys.stderr,
        )
        return 1

    if resp.status_code >= 400:
        try:
            data = resp.json()
            err = data.get("error", "")
            if err:
                print(str(err), file=sys.stderr)
                return 1
            if "message" in data:
                print(data["message"], file=sys.stderr)
                return 1
        except Exception:
            pass
        print(resp.text, file=sys.stderr)
        return 1

    content_type = resp.headers.get("Content-Type", "")
    if "application/json" in content_type:
        try:
            data = resp.json()
        except Exception:
            print(resp.text)
            return 0 if resp.ok else 1

        err = data.get("error", "")
        res = data.get("result", None)
        if err:
            print(err, file=sys.stderr)
            if res is not None:
                print(json.dumps(res, indent=2, ensure_ascii=False), file=sys.stderr)
            return 1

        if isinstance(res, str):
            print(res, end="" if res.endswith("\n") else "\n")
        else:
            if pretty_print:
                print(json.dumps(res, indent=2, ensure_ascii=False))
            else:
                print(json.dumps(res, ensure_ascii=False))
        return 0 if resp.ok else 1

    if resp.text:
        try:
            value = json.loads(resp.text)
            if pretty_print:
                print(json.dumps(value, indent=2, ensure_ascii=False))
            else:
                print(json.dumps(value, ensure_ascii=False))
            return 0 if resp.ok else 1
        except Exception:
            print(resp.text, end="" if resp.text.endswith("\n") else "\n")

    return 0 if resp.ok else 1


def run_login(base_url: Optional[str] = None) -> int:
    try:
        interactive_login(default_base_url=base_url)
        print("Login successful. Configuration saved.")
        return 0
    except RuntimeError as exc:
        print(str(exc), file=sys.stderr)
        return 2
    except requests.RequestException:
        print(
            "Cannot reach the platform. Is the platform running?\n"
            "If the URL is wrong, run 'datatailr login' to update it.",
            file=sys.stderr,
        )
        return 2


def _ensure_login(
    force_login: bool, base_url: Optional[str] = None
) -> Tuple[Config, bool]:
    if not force_login:
        try:
            cfg = Config.load()
            validate_remote_session(cfg.base_url, cfg.jwt_cookie, timeout_s=10)
            return cfg, False
        except Exception:
            pass

    print("Please login to your Datatailr installation.")
    cfg = interactive_login(default_base_url=base_url)
    return cfg, True


def run_setup_ssh(
    workstation: Optional[str] = None,
    force_login: bool = False,
    base_url: Optional[str] = None,
) -> int:
    try:
        cfg, _ = _ensure_login(force_login=bool(force_login), base_url=base_url)
        api_json = load_packaged_api_json()
        selected_workstation = workstation
        if not selected_workstation:
            username = username_from_jwt(cfg.jwt_cookie)
            selected_workstation = select_workstation_interactively(
                cfg, api_json, username, _api_call
            )

        host_alias = setup_ssh_config(cfg, selected_workstation)
        message = "Connection configured! You can connect with:"
        if sys.stdout.isatty():
            print(f"\033[4m{message}\033[0m")
        else:
            print(message)
        command_line = f"ssh {host_alias}"
        if sys.stdout.isatty():
            print(f"\033[1m{command_line}\033[0m")
        else:
            print(command_line)
        return 0
    except RuntimeError as exc:
        print(str(exc), file=sys.stderr)
        return 2
    except requests.RequestException as exc:
        print(str(exc), file=sys.stderr)
        return 2


def main(argv: Optional[List[str]] = None) -> int:
    argv = argv if argv is not None else sys.argv[1:]
    parser = build_parser()

    if argv and argv[0] in {"login", "setup-ssh", "setup-cli"}:
        print(
            f"'{argv[0]}' is a local setup command. Use 'datatailr {argv[0]}' instead.",
            file=sys.stderr,
        )
        return 2

    if not argv:
        try:
            cfg = Config.load()
            return send_help(cfg, None, None)
        except Exception:
            parser.print_help()
            return 1

    cfg = Config.load()
    api_json = load_packaged_api_json()
    try:
        return send_request(cfg, api_json, argv)
    except RuntimeError as exc:
        print(str(exc), file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())
