"""
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 setup_ssh_config
from datatailr._remote_client.utils import (
    Config,
    build_headers,
    fetch_openapi,
    interactive_login,
    is_auth_failure,
    load_cached_openapi,
    username_from_jwt,
)


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

    for param in operation.get("parameters", []) or []:
        name = param.get("name")
        if not name:
            continue
        schema = param.get("schema") or {}
        param_type = schema.get("type")
        long_requires[name] = param_type != "boolean"
        short_name = param.get("x-short-name")
        if short_name:
            short_to_long[short_name] = name

    rb = operation.get("requestBody", {}) or {}
    content = (rb.get("content") or {}).get("application/json") or {}
    schema = content.get("schema") or {}
    for name, prop in (schema.get("properties") or {}).items():
        prop_type = (prop or {}).get("type")
        long_requires[name] = prop_type != "boolean"
        short_name = (prop or {}).get("x-short-name")
        if short_name:
            short_to_long[short_name] = name

    return short_to_long, long_requires


def parse_options_from_args(
    argv_tail: List[str],
    operation: Dict[str, Any],
    short_to_long: Dict[str, str],
    long_requires: Dict[str, bool],
) -> Tuple[List[str], Dict[str, Any]]:
    valid_long = {p.get("name") for p in operation.get("parameters", []) or []}
    rb = operation.get("requestBody", {}) or {}
    content = (rb.get("content") or {}).get("application/json") or {}
    schema = content.get("schema") or {}
    valid_long.update((schema.get("properties") or {}).keys())
    valid_long.discard(None)

    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:]
            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 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)
    )
    return {
        "repo_archive_b64": archive_b64,
        "repo_dir_name": repo_dir_name,
        "repo_module_path": module_path,
    }


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_operation_for_path(
    spec: Dict[str, Any], api_path: str
) -> Tuple[str, Dict[str, Any]]:
    paths = spec.get("paths", {})
    if api_path not in paths:
        raise RuntimeError(f"Unknown command: no OpenAPI path '{api_path}'")

    methods_obj = paths[api_path]
    candidates = [
        (m.lower(), op)
        for m, op in methods_obj.items()
        if m.lower() in {"get", "post", "put", "patch", "delete"}
    ]
    if not candidates:
        raise RuntimeError(f"No HTTP operation found under '{api_path}'")
    if len(candidates) != 1:
        raise RuntimeError(
            f"Expected exactly 1 method for '{api_path}', found: {[m for m, _ in candidates]}"
        )
    return candidates[0][0].upper(), candidates[0][1]


def build_request_from_openapi(
    operation: Dict[str, Any],
    argv_tail: List[str],
    extra_query: Optional[Dict[str, Any]] = None,
    extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    query: Dict[str, Any] = dict(extra_query or {})
    body: Dict[str, Any] = dict(extra_body or {})

    all_query_params: List[Dict[str, Any]] = []
    for param in operation.get("parameters", []) or []:
        if param.get("in") == "query":
            all_query_params.append(param)

    required_query_names = [
        param["name"] for param in all_query_params if param.get("required") is True
    ]
    optional_query_names = [
        param["name"] for param in all_query_params if param.get("required") is not True
    ]

    required_body_names: List[str] = []
    optional_body_names: List[str] = []

    rb = operation.get("requestBody", {})
    if rb:
        content = (rb.get("content") or {}).get("application/json")
        if content:
            schema = content.get("schema") or {}
            required_body_names = list(schema.get("required") or [])
            props = schema.get("properties", {})
            for key in props.keys():
                if key not in required_body_names:
                    optional_body_names.append(key)

    required_query_names = [name for name in required_query_names if name not in query]
    required_body_names = [name for name in required_body_names if name not in body]
    optional_query_names = [name for name in optional_query_names if name not in query]
    optional_body_names = [name for name in optional_body_names if name not in body]

    min_expected = len(required_query_names) + len(required_body_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)}."
        )

    cursor = 0
    for name in required_query_names:
        query[name] = argv_tail[cursor]
        cursor += 1

    for name in required_body_names:
        value = argv_tail[cursor]
        cursor += 1
        if name == "data" and is_local_file_arg(value):
            job_file = Path(value)
            job_json = encode_file_as_json(value)
            extra_bundle = try_to_prepare_repo_bundle(job_json, job_file)
            body[name] = job_json
            body.update(extra_bundle)
        else:
            body[name] = value

    for name in optional_query_names:
        if cursor >= len(argv_tail):
            break
        query[name] = argv_tail[cursor]
        cursor += 1

    for name in optional_body_names:
        if cursor >= len(argv_tail):
            break
        body[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, spec: Dict[str, Any], full_argv: List[str], timeout_s: int = 60
) -> int:
    def handle_unknown_command(target_cmd: str) -> int:
        cmd_exists = any(
            path.startswith(f"/api/{target_cmd}/") for path in spec.get("paths", {})
        )
        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}"
    try:
        method, operation = get_operation_for_path(spec, api_path)
    except RuntimeError:
        return handle_unknown_command(cmd)

    if cmd == "blob" and subcmd == "cp":
        print("dt blob cp is not yet supported in remote CLI.", file=sys.stderr)
        return 2

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

    try:
        if method == "GET":
            query, body = build_request_from_openapi(
                operation, positionals, extra_query=options
            )
        else:
            query, body = build_request_from_openapi(
                operation, positionals, extra_body=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

    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,
    spec: Dict[str, Any],
    cmd: str,
    subcmd: str,
    argv_tail: List[str],
    timeout_s: int = 60,
) -> Any:
    api_path = f"/api/{cmd}/{subcmd}"
    method, operation = get_operation_for_path(spec, api_path)

    short_to_long, long_requires = load_api_option_metadata(operation)
    positionals, options = parse_options_from_args(
        argv_tail, operation, short_to_long, long_requires
    )
    if method == "GET":
        query, body = build_request_from_openapi(
            operation, positionals, extra_query=options
        )
    else:
        query, body = build_request_from_openapi(
            operation, positionals, extra_body=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 _job_state_from_get_result(result: Any) -> Optional[str]:
    if isinstance(result, dict):
        state = result.get("state")
        if isinstance(state, str) and state:
            return state.lower()
        for key in ("jobs", "items", "data"):
            nested = result.get(key)
            if isinstance(nested, list) and nested:
                first = nested[0]
                if isinstance(first, dict):
                    nested_state = first.get("state")
                    if isinstance(nested_state, str) and nested_state:
                        return nested_state.lower()
        return None

    if isinstance(result, list) and result:
        first = result[0]
        if isinstance(first, dict):
            state = first.get("state")
            if isinstance(state, str) and state:
                return state.lower()

    return None


def _job_rows_from_ls_result(result: Any) -> List[Dict[str, Any]]:
    if isinstance(result, str):
        text = result.strip()
        if not text:
            return []
        try:
            parsed = json.loads(text)
        except Exception:
            return []
        return _job_rows_from_ls_result(parsed)

    if isinstance(result, list):
        return [row for row in result if isinstance(row, dict)]
    if isinstance(result, dict):
        for key in ("jobs", "items", "data", "result"):
            rows = result.get(key)
            if isinstance(rows, list):
                return [row for row in rows if isinstance(row, dict)]
    return []


def _format_memory_bytes(value: Any) -> str:
    if isinstance(value, (int, float)):
        memory = float(value)
    elif isinstance(value, str) and value.strip():
        try:
            memory = float(value.strip())
        except ValueError:
            return value
    else:
        return "unknown"

    units = ["B", "KiB", "MiB", "GiB", "TiB"]
    unit_index = 0
    while memory >= 1024 and unit_index < len(units) - 1:
        memory /= 1024.0
        unit_index += 1
    if unit_index == 0:
        return f"{int(memory)} {units[unit_index]}"
    return f"{memory:.2f} {units[unit_index]}"


def _normalize_requirements(value: Any) -> List[str]:
    if isinstance(value, list):
        return [str(item).strip() for item in value if str(item).strip()]
    if not isinstance(value, str):
        return []

    text = value.strip()
    if not text:
        return []

    parts: List[str] = []
    for line in text.splitlines():
        for token in line.split(","):
            item = token.strip()
            if item:
                parts.append(item)
    return parts


def _format_requirements_display(requirements: List[str], max_length: int = 120) -> str:
    if not requirements:
        return "none"

    joined = ", ".join(requirements)
    if len(joined) <= max_length:
        return joined

    if max_length <= 3:
        return "..."

    return joined[: max_length - 3].rstrip() + "..."


def _list_workspace_candidates(
    cfg: Config, spec: Dict[str, Any], username: str
) -> List[Dict[str, Any]]:
    workspace_names: Dict[str, None] = {}
    suffix = f"-{username}"

    result = _api_call(
        cfg,
        spec,
        "job",
        "ls",
        ["-j", "--filter", "type=workspace,environment=dev"],
        timeout_s=30,
    )
    rows = _job_rows_from_ls_result(result)
    for row in rows:
        job_name = row.get("name")
        if not isinstance(job_name, str) or not job_name:
            continue

        if job_name.endswith(suffix) and len(job_name) > len(suffix):
            workspace_name = job_name[: -len(suffix)]
        else:
            workspace_name = job_name

        workspace_names[workspace_name] = None

    candidates: List[Dict[str, Any]] = []
    for workspace_name in sorted(workspace_names):
        details: Dict[str, Any] = {
            "workspace_name": workspace_name,
            "display_name": workspace_name,
            "cpu": None,
            "memory": None,
            "python_requirements": [],
        }

        try:
            result = _api_call(
                cfg,
                spec,
                "job",
                "get",
                [workspace_name, "-j", "-p", "-e", "dev"],
                timeout_s=30,
            )
            rows = _job_rows_from_ls_result(result)
            if rows:
                row = rows[0]
                if row.get("per_user_job") is not True:
                    continue

                image = row.get("image")
                requirements_value = None
                if isinstance(image, dict):
                    requirements_value = image.get("python_requirements")

                display_name = row.get("display_name")
                details["display_name"] = (
                    display_name
                    if isinstance(display_name, str) and display_name
                    else workspace_name
                )
                details["cpu"] = row.get("num_cpus")
                details["memory"] = row.get("memory")
                details["python_requirements"] = _normalize_requirements(
                    requirements_value
                )
        except Exception:
            # Keep list usable even if metadata fetch fails for a workspace.
            pass

        candidates.append(details)

    return candidates


def _select_workspace_interactively(
    cfg: Config, spec: Dict[str, Any], username: str
) -> str:
    workspaces = _list_workspace_candidates(cfg, spec, username)
    if not workspaces:
        raise RuntimeError(
            "Could not find any workspace jobs in dev for your account. "
            "Pass workspace explicitly: 'datatailr setup-ssh <workspace>'."
        )

    choose_label = "Choose a workspace:"
    if sys.stdout.isatty():
        print(f"\033[4m{choose_label}\033[0m")
    else:
        print(choose_label)
    for index, workspace in enumerate(workspaces, start=1):
        display_name = workspace["display_name"]
        name = workspace["workspace_name"]
        cpu = workspace.get("cpu")
        cpu_display = str(cpu) if cpu is not None else "unknown"
        memory_display = _format_memory_bytes(workspace.get("memory"))
        reqs = workspace.get("python_requirements") or []
        reqs_display = _format_requirements_display(reqs)

        print(
            f"  {index}. {display_name} ({name}) | "
            f"cpu={cpu_display} Core | memory={memory_display} | "
            f"python_requirements=[{reqs_display}]"
        )

    while True:
        raw = input(f"Select workspace [1-{len(workspaces)}], c to cancel: ").strip()
        if not raw:
            continue
        if raw.lower() in {"c", "cancel", "q", "quit", "exit"}:
            raise RuntimeError("Workspace selection cancelled by user.")
        if raw.isdigit():
            choice = int(raw)
            if 1 <= choice <= len(workspaces):
                selected = str(workspaces[choice - 1]["workspace_name"])
                return selected
        print("Invalid selection. Enter a valid number from the list.")


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 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() -> int:
    try:
        interactive_login()
        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) -> Tuple[Config, bool]:
    if not force_login:
        try:
            cfg = Config.load()
            # Validate that token/base_url still work
            fetch_openapi(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()
    return cfg, True


def run_setup_ssh(workspace: Optional[str] = None, force_login: bool = False) -> int:
    try:
        cfg, _ = _ensure_login(force_login=bool(force_login))
        spec = fetch_openapi(cfg.base_url, cfg.jwt_cookie, timeout_s=15)
        selected_workspace = workspace
        if not selected_workspace:
            username = username_from_jwt(cfg.jwt_cookie)
            selected_workspace = _select_workspace_interactively(cfg, spec, username)

        host_alias = setup_ssh_config(cfg, selected_workspace)
        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()
    spec = load_cached_openapi()
    try:
        return send_request(cfg, spec, argv)
    except RuntimeError as exc:
        print(str(exc), file=sys.stderr)
        return 2


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