# *************************************************************************
#
#  Copyright (c) 2026 - Datatailr Inc.
#  All Rights Reserved.
#
#  This file is part of Datatailr and subject to the terms and conditions
#  defined in 'LICENSE.txt'. Unauthorized copying and/or distribution
#  of this file, in parts or full, via any medium is strictly prohibited.
# *************************************************************************

from __future__ import annotations

import base64
import getpass
import json
import os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional

import requests

CONFIG_DIR = Path.home() / ".dt" / "remote_client"
CONFIG_PATH = CONFIG_DIR / "remote_client.cfg"
OPENAPI_CACHE_PATH = CONFIG_DIR / "openapi.json"
COOKIE_NAME = "X-Datatailr-Oidc-Data"


@dataclass
class Config:
    base_url: str
    jwt_cookie: str

    @staticmethod
    def load() -> "Config":
        if not CONFIG_PATH.exists():
            raise RuntimeError(
                "No config found for remote cli - please run 'datatailr login'."
            )
        data = json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
        base_url = data.get("base_url")
        jwt_cookie = data.get("jwt_cookie")
        if not base_url or not jwt_cookie:
            raise RuntimeError(
                "Invalid config for remote cli - please run 'datatailr login'."
            )
        return Config(base_url=base_url, jwt_cookie=jwt_cookie)

    def save(self) -> None:
        CONFIG_DIR.mkdir(parents=True, exist_ok=True)
        CONFIG_PATH.write_text(
            json.dumps(
                {"base_url": self.base_url, "jwt_cookie": self.jwt_cookie},
                indent=2,
            ),
            encoding="utf-8",
        )
        try:
            os.chmod(CONFIG_PATH, 0o600)
        except Exception:
            pass


def normalize_base_url(base_url: str) -> str:
    url = base_url.strip().rstrip("/")
    if not url.startswith("http://") and not url.startswith("https://"):
        url = "http://" + url if "localhost" in url else "https://" + url
    return url


def fetch_openapi(
    base_url: str, jwt_cookie: str, timeout_s: int = 10
) -> Dict[str, Any]:
    url = f"{base_url}/api/openapi.json"
    for attempt in range(3):
        resp = requests.get(
            url,
            headers={"Cookie": f"{COOKIE_NAME}={jwt_cookie}"},
            timeout=timeout_s,
        )
        if resp.status_code >= 400:
            raise RuntimeError(
                f"Failed to fetch OpenAPI: {resp.status_code}\n{resp.text[:400]}"
            )
        try:
            return resp.json()
        except requests.exceptions.JSONDecodeError:
            if attempt < 2:
                time.sleep(0.3 * (attempt + 1))
                continue
            break

    raise RuntimeError(
        "Failed to parse OpenAPI response (not JSON). If datatailr was just started, it might take around a minute to be ready.\n"
    )


def cache_openapi(spec: Dict[str, Any]) -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    OPENAPI_CACHE_PATH.write_text(json.dumps(spec, indent=2), encoding="utf-8")


def load_cached_openapi() -> Dict[str, Any]:
    if not OPENAPI_CACHE_PATH.exists():
        raise RuntimeError(
            f"No cached OpenAPI at {OPENAPI_CACHE_PATH}. Please run 'datatailr login'."
        )
    return json.loads(OPENAPI_CACHE_PATH.read_text(encoding="utf-8"))


def auth_login(base_url: str, username: str, password: str, timeout_s: int = 30) -> str:
    sess = requests.Session()
    resp = sess.post(
        f"{base_url}/login",
        params={"response_type": "id_token"},
        auth=(username, password),
        timeout=timeout_s,
    )
    if resp.status_code == 401:
        raise RuntimeError("Login failed: invalid username or password.")
    if resp.status_code >= 400:
        raise RuntimeError(f"Login failed: HTTP {resp.status_code}\n{resp.text[:500]}")

    jwt = sess.cookies.get(COOKIE_NAME)
    if not jwt:
        raise RuntimeError("Login succeeded but JWT token is not found.")
    return jwt


def interactive_login(default_base_url: Optional[str] = None) -> Config:
    base_default = default_base_url
    if base_default is None:
        base_default = "http://localhost"
        try:
            base_default = Config.load().base_url
        except Exception:
            pass

    base_url_in = input(f"Base URL (default {base_default}): ").strip() or base_default
    base_url = normalize_base_url(base_url_in)
    username = input("Username: ").strip()
    password = getpass.getpass("Password: ")

    jwt = auth_login(base_url, username, password)
    spec = fetch_openapi(base_url, jwt)
    cache_openapi(spec)

    cfg = Config(base_url=base_url, jwt_cookie=jwt)
    cfg.save()
    return cfg


def build_headers(cfg: Config) -> Dict[str, str]:
    return {
        "Cookie": f"{COOKIE_NAME}={cfg.jwt_cookie}",
        "Accept": "text/plain, application/json;q=0.9, */*;q=0.8",
        "User-Agent": "dt-remote-cli/0.1",
    }


def is_auth_failure(resp: requests.Response) -> bool:
    if resp.status_code == 401:
        return True
    if resp.status_code in {301, 302, 303, 307, 308} and "login" in resp.url:
        return True
    content_type = resp.headers.get("Content-Type", "")
    if "text/html" in content_type.lower():
        body = resp.text or ""
        if "Enter your login credentials" in body or (
            "<title" in body and "Login" in body
        ):
            return True
    return False


def jwt_payload(jwt_cookie: str) -> Dict[str, Any]:
    parts = jwt_cookie.split(".")
    if len(parts) < 2:
        raise RuntimeError("Invalid auth token in remote CLI config.")

    payload = parts[1]
    payload += "=" * ((4 - len(payload) % 4) % 4)
    decoded = base64.urlsafe_b64decode(payload.encode("utf-8")).decode("utf-8")
    data = json.loads(decoded)
    if not isinstance(data, dict):
        raise RuntimeError("Invalid auth token payload.")
    return data


def username_from_jwt(jwt_cookie: str) -> str:
    payload = jwt_payload(jwt_cookie)
    username = payload.get("preferred_username")
    if not isinstance(username, str) or not username:
        raise RuntimeError("Could not resolve username from auth token.")
    return username
