import collections
import datetime
import json
import logging
import os
import re
import subprocess
import tempfile
import textwrap
import threading
import time
from pathlib import Path

import urllib3

from ... import util

logger = logging.getLogger("atex.provisioner.testingfarm")

DEFAULT_API_URL = "https://api.testing-farm.io"

DEFAULT_RESERVE_TEST = {
    "url": "https://github.com/RHSecurityCompliance/atex-reserve",
    "ref": "0.12",
    "path": ".",
    "name": "/plans/reserve",
}

# final states of a request,
# https://gitlab.com/testing-farm/nucleus/-/blob/main/api/src/tft/nucleus/api/core/schemes/test_request.py
END_STATES = ("error", "complete", "canceled")

# always have at most 10 outstanding HTTP requests to every given API host,
# shared by all instances of all classes here, to avoid flooding the host
# by multi-threaded users
_http = urllib3.PoolManager(
    maxsize=10,
    block=True,
    retries=urllib3.Retry(
        total=24,
        # account for API restarts / short outages
        backoff_factor=10,
        backoff_max=3600,
        # retry on API server errors too, not just connection issues
        status=10,
        status_forcelist={403,404,408,429,500,502,503,504},
        # retry POST as well, even if risky
        allowed_methods=urllib3.Retry.DEFAULT_ALLOWED_METHODS | {"POST"},
    ),
)


class TestingFarmError(Exception):
    def __init__(self, message, reply=None):
        super().__init__(message)
        self.reply = reply


class APIError(TestingFarmError):
    pass


# TODO docstrings for these:

class BadHTTPError(TestingFarmError):
    pass


class GoneAwayError(TestingFarmError):
    pass


class TestingFarmAPI:
    """
    A python interface for the Testing Farm HTTP API, closely matching
    functionality provided by it, returning python dictionaries that
    correspond to JSON replies of the HTTP API functions.
    """

    def __init__(self, url=DEFAULT_API_URL, token=None):
        """
        - `url` is Testing Farm API URL, a sensible default is used
          if unspecified.

        - `token` is a secret API token generated by Testing Farm admins,
          if empty, the `TESTING_FARM_API_TOKEN` env var is read instead.

        Note that token-less operation is supported, with limited functionality.
        """
        self.api_url = url
        self.api_token = token or os.environ.get("TESTING_FARM_API_TOKEN")

    def _query(self, method, path, *args, headers=None, auth=True, version="v0.1", **kwargs):
        url = f"{self.api_url}/{version}{path}"
        if self.api_token and auth:
            if headers is not None:
                headers["Authorization"] = f"Bearer {self.api_token}"
            else:
                headers = {"Authorization": f"Bearer {self.api_token}"}

        reply = _http.request(method, url, *args, headers=headers, preload_content=False, **kwargs)

        if reply.status != 200 and not reply.data:
            raise APIError(f"got HTTP {reply.status} on {method} {url}", reply)

        if reply.headers.get("Content-Type") != "application/json":
            raise BadHTTPError(
                f"HTTP {reply.status} on {method} {url} is not application/json",
                reply,
            )

        try:
            decoded = reply.json()
        except json.decoder.JSONDecodeError:
            raise BadHTTPError(
                f"failed to decode JSON for {method} {url}: {reply.data}",
                reply,
            ) from None

        if reply.status != 200:
            raise APIError(f"got HTTP {reply.status} on {method} {url}: {decoded}", reply)

        return decoded

    def whoami(self):
        if not self.api_token:
            raise ValueError("whoami() requires an auth token")
        if hasattr(self, "_whoami_cached"):
            return self._whoami_cached
        else:
            self._whoami_cached = self._query("GET", "/whoami")
            return self._whoami_cached

    def about(self):
        return self._query("GET", "/about")

    def composes(self, ranch=None):
        """
        - `ranch` is `public` or `redhat`, autodetected if token was given.
        """
        if not ranch:
            if not self.api_token:
                raise ValueError("composes() requires an auth token to identify ranch")
            ranch = self.whoami()["token"]["ranch"]
        return self._query("GET", f"/composes/{ranch}", version="v0.2")

    def search_requests(
        self, *, state, ranch=None,
        mine=True, user_id=None, token_id=None,
        created_before=None, created_after=None,
    ):
        """
        - `state` is one of `running`, `queued`, etc., and is required by the
          API.

        - `ranch` is `public` or `redhat`, or (probably?) all if left empty.

        - If `mine` is `True` and a token was given, return only requests for
          that token (user), otherwise return *all* requests (use extra filters
          pls).

        - `user_id` and `token_id` are search API parameters - if not given and
          `mine` is `True`, these are extracted from a user-provided token.

        - `created_*` take ISO 8601 formatted strings, as returned by the API
          elsewhere, ie. `YYYY-MM-DD` or `YYYY-MM-DDTHH:MM:SS` (or with `.MS`),
          without timezone (UTC is used always).
        """
        fields = {"state": state}
        if ranch:
            fields["ranch"] = ranch
        if created_before:
            fields["created_before"] = created_before
        if created_after:
            fields["created_after"] = created_after

        if user_id or token_id:
            if user_id:
                fields["user_id"] = user_id
            if token_id:
                fields["token_id"] = token_id
        elif mine:
            if not self.api_token:
                raise ValueError("search_requests(mine=True) requires an auth token")
            fields["token_id"] = self.whoami()["token"]["id"]
            fields["user_id"] = self.whoami()["user"]["id"]

        return self._query("GET", "/requests", fields=fields, auth=mine)

    def search_requests_paged(self, *args, page=43200, **kwargs):
        """
        An unofficial wrapper for search_requests() that can search a large
        interval incrementally (in "pages") and yield batches of results.

        Needs `created_after`, with `created_before` defaulting to `now()`.

        - `page` specifies the time interval of one page, in seconds.

        - `args` and `kwargs` are passed to `search_requests()`.
        """
        assert "created_after" in kwargs, "at least 'created_after' is needed for paging"

        def from_iso8601(date):
            dt = datetime.datetime.fromisoformat(date)
            # if no TZ is specified, treat it as UTC, not localtime
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=datetime.UTC)
            # convert to UTC
            else:
                dt = dt.astimezone(datetime.UTC)
            return dt

        after = from_iso8601(kwargs["created_after"])
        if kwargs.get("created_before"):
            before = from_iso8601(kwargs["created_before"])
        else:
            before = datetime.datetime.now(datetime.UTC)

        # scale down page size to fit between after/before
        page = min(page, (before - after).total_seconds())

        start = after
        while start < before:
            end = start + datetime.timedelta(seconds=page)
            # clamp to real 'before'
            end = min(end, before)
            new_kwargs = kwargs | {
                "created_after": start.isoformat(),
                "created_before": end.isoformat(),
            }
            found = self.search_requests(*args, **new_kwargs)
            if found is not None:
                yield from found
            start = end

    def get_request(self, request_id):
        """
        - `request_id` is the UUID (string) of the request.
        """
        return self._query("GET", f"/requests/{request_id}")

    def submit_request(self, spec):
        """
        - `spec` is a big dictionary with 'test', 'environment', 'settings', etc.
        keys that specify what should be run and where.
        """
        if not self.api_token:
            raise ValueError("submit_request() requires an auth token")
        return self._query("POST", "/requests", json=spec)

    def cancel_request(self, request_id):
        """
        - `request_id` is the UUID (string) of the request.
        """
        return self._query("DELETE", f"/requests/{request_id}")


class Request:
    """
    A higher-level API for submitting, querying, and cancelling a Testing Farm
    request.
    """

    # actually query the TestingFarm API at most every X seconds,
    # re-using cached state between updates
    api_query_limit = 30

    def __init__(self, id=None, api=None, initial_data=None):
        """
        - `id` is a Testing Farm request UUID.

        - `api` is a TestingFarmAPI instance - if unspecified, a new one
          is instantiated.

        - `initial_data` (dict) can be used to pre-fill an initial Request
          state.
        """
        self.id = id
        self.api = api or TestingFarmAPI()
        self.data = initial_data or {}
        self.next_query = 0

    def submit(self, spec):
        """
        - `spec` is a big dictionary with 'test', 'environment', 'settings',
          etc. keys that specify what should be run and where.
        """
        if self.id:
            raise ValueError("this Request instance already has 'id', refusing submit")
        self.data = self.api.submit_request(spec)
        self.id = self.data["id"]

    def _refresh(self):
        if not self.id:
            return
        if time.monotonic() > self.next_query:
            self.data = self.api.get_request(self.id)
            self.next_query = time.monotonic() + self.api_query_limit

    def cancel(self):
        if not self.id:
            return
        data = self.api.cancel_request(self.id)
        self.id = None
        self.data = {}
        return data

    def alive(self):
        if not self.id:
            return False
        self._refresh()
        return self.data["state"] not in END_STATES

    def assert_alive(self):
        if not self.alive():
            state = self.data["state"]
            raise GoneAwayError(f"request {self.id} not alive anymore, entered: {state}")

    def wait_for_state(self, state):
        """
        - `state` is a string or a tuple of states to wait for.
        """
        watched = (state,) if isinstance(state, str) else state
        while True:
            self._refresh()
            if self.data["state"] in watched:
                break
            # if the request ended in one of END_STATES and the above condition
            # did not catch it, the wait will never end
            if self.data["state"] in END_STATES:
                raise GoneAwayError(f"request {self.id} ended with {self.data['state']}")

    def __repr__(self):
        return f"Request(id={self.id})"

    def __str__(self):
        self._refresh()
        # python has no better dict-pretty-printing logic
        return json.dumps(self.data, sort_keys=True, indent=4)

    def __contains__(self, item):
        self._refresh()
        return item in self.data

    def __getitem__(self, key):
        self._refresh()
        return self.data[key]


class PipelineLogStreamer:
    """
    Line buffer for querying Testing Farm pipeline.log using HTTP Range header
    to "stream" its contents over time (over many requests), never having to
    re-read old pipeline.log content.
    """

    # how frequently to check for pipeline.log updates (seconds)
    pipeline_query_limit = 30

    def __init__(self, request):
        self.request = request

    def _wait_for_entry(self):
        while True:
            self.request.wait_for_state("running")

            try:
                if "run" not in self.request or "artifacts" not in self.request["run"]:
                    continue

                artifacts = self.request["run"]["artifacts"]
                if not artifacts:
                    continue

                log = f"{artifacts}/pipeline.log"
                reply = _http.request("HEAD", log)
                # 404: TF has a race condition of adding the .log entry without
                #      it being created
                # 403: happens on internal OSCI artifacts server, probably
                #      due to similar reasons (folder exists without log)
                if reply.status in (404,403):
                    logger.info(f"got {reply.status} for {log}, retrying")
                    continue
                elif reply.status != 200:
                    raise APIError(f"got HTTP {reply.status} on HEAD {log}", reply)

                logger.info(f"artifacts: {artifacts}")

                return log

            finally:
                time.sleep(self.pipeline_query_limit)

    def __iter__(self):
        url = self._wait_for_entry()
        buffer = ""
        bytes_read = 0
        while True:
            self.request.assert_alive()

            try:
                headers = {"Range": f"bytes={bytes_read}-"}
                # load all returned data via .decode() rather than streaming it
                # in chunks, because we don't want to leave the connection open
                # (blocking others) while the user code runs between __next__ calls
                reply = _http.request("GET", url, headers=headers)

                # 416=Range Not Satisfiable, typically meaning "no new data to send"
                if reply.status == 416:
                    continue
                # 200=OK or 206=Partial Content
                elif reply.status not in (200,206):
                    raise BadHTTPError(f"got {reply.status} when trying to GET {url}", reply)

                bytes_read += len(reply.data)
                buffer += reply.data.decode(errors="ignore")

                while (index := buffer.find("\n")) != -1:
                    yield buffer[:index]
                    buffer = buffer[index+1:]

            finally:
                time.sleep(self.pipeline_query_limit)


class Reserve:
    r"""
    An abstraction for (ab)using Testing Farm for OS reservations, by submitting
    a dummy test that sets up user-provided SSH key access and enters a no-op
    state, allowing ad-hoc access/use by the user.

    When used in a context manager, it produces a ReservedMachine tuple with
    connection details for an ssh client:

        with Reserve(compose="CentOS-Stream-9", timeout=720) as m:
            subprocess.run(["ssh", "-i", m.ssh_key, f"{m.user}@{m.host}", "ls /"])
    """

    Reserved = collections.namedtuple(
        "ReservedMachine",
        ("host", "port", "user", "ssh_key", "request"),
    )

    def __init__(
        self, *, compose, arch="x86_64", pool=None, hardware=None, kickstart=None,
        timeout=60, ssh_key=None, source_host=None,
        reserve_test=None, variables=None, secrets=None, tags=None,
        api=None,
    ):
        """
        - `compose` (str) is the OS to install, chosen from the composes
          supported by the Testing Farm ranch of the authenticated user.

        - `arch` (str) is one of 'x86_64', 's390x', etc.

        - `pool` (str) is a name of a Testing Farm infrastructure pool.

        - `hardware` (dict) is a complex specification of hardware properties
          the reserved system should have, see:
          https://docs.testing-farm.io/Testing%20Farm/0.1/test-request.html#hardware

        - `kickstart` (dict) is a Beaker-style specification of Anaconda
          Kickstart hacks, passed directly to Testing Farm POST /requests API.

        - `timeout` (int) is the maximum time **in minutes** a Testing Farm
          request is alive, which includes initial creation, waiting in queue,
          preparing an OS, and the entire reservation period.

          Make sure to set it high enough (not just the pure reservation time).

        - `ssh_key` (str) is a path to an OpenSSH private key file (with an
          associated public key file in .pub), to be added to the reserved OS.

          If unspecified, an attempt to read `~/.ssh/id_rsa` will be made, and
          if that is also unsuccessful, a temporary keypair will be generated.

        - `source_host` (str) is an IPv4 network specified as ie. `1.2.3.4/32`
          to be allowed incoming traffic to the reserved system (such as ssh).

          If unspecified, an Internet service will be queried to get an outside-
          facing address of the current system.

          Ignored on the `redhat` ranch.

        - `reserve_test` is a dict with a fmf test specification to be run on
          the target system to reserve it, ie.:
              {
                  "url": "https://some-host/path/to/repo",
                  "ref": "main",
                  "name": "/plans/reserve",
              }

        - `variables` and `secrets` are dicts with environment variable
          key/values exported for the reserve test - variables are visible via
          TF API, secrets are not (but can still be extracted from pipeline
          log).

        - `tags` is a dict of custom key/values to be submitted in TF Request as
          `environments->settings->provisioning->tags`, useful for storing
          custom metadata to be queried later.

        - `api` is a TestingFarmAPI instance - if unspecified, a new one
          is instantiated.
        """
        spec = {
            "test": {
                "fmf": reserve_test or DEFAULT_RESERVE_TEST,
            },
            "environments": [{
                "arch": arch,
                "os": {
                    "compose": compose,
                },
                "settings": {
                    "pipeline": {
                        "skip_guest_setup": True,
                    },
                    "provisioning": {
                        "tags": {
                            "ArtemisUseSpot": "false",
                        },
                    },
                },
            }],
            "settings": {
                "pipeline": {
                    "timeout": timeout,
                },
            },
        }
        spec_env = spec["environments"][0]
        if pool:
            spec_env["pool"] = pool
        if hardware:
            spec_env["hardware"] = hardware
        if kickstart:
            spec_env["kickstart"] = kickstart
        if variables:
            spec_env["variables"] = variables
        spec_env["secrets"] = secrets.copy() if secrets else {}  # we need it for ssh pubkey
        if tags:
            spec_env["settings"]["provisioning"]["tags"] |= tags

        self._spec = spec
        self._ssh_key = Path(ssh_key) if ssh_key else None
        self._source_host = source_host
        self.api = api or TestingFarmAPI()

        self.lock = threading.RLock()
        self.request = None
        self._tmpdir = None

    @staticmethod
    def _guess_host_ipv4():
        curl_agent = {"User-Agent": "curl/1.2.3"}
        try:
            r = _http.request("GET", "https://ifconfig.me", headers=curl_agent)
            if r.status != 200:
                raise ConnectionError
        except (ConnectionError, urllib3.exceptions.RequestError):
            r = _http.request("GET", "https://ifconfig.co", headers=curl_agent)
        return r.data.decode().strip()

    def reserve(self):
        with self.lock:
            if self.request:
                raise RuntimeError("reservation already in progress")

        spec = self._spec.copy()
        spec_env = spec["environments"][0]

        # add source_host firewall filter on the public ranch
        if self.api.whoami()["token"]["ranch"] == "public":
            source_host = self._source_host or f"{self._guess_host_ipv4()}/32"
            ingress_rule = {
                "type": "ingress",
                "protocol": "-1",
                "cidr": source_host,
                "port_min": 0,
                "port_max": 65535,
            }
            provisioning = spec_env["settings"]["provisioning"]
            if "security_group_rules_ingress" in provisioning:
                provisioning["security_group_rules_ingress"].append(ingress_rule)
            else:
                provisioning["security_group_rules_ingress"] = [ingress_rule]

        try:
            # read user-provided ssh key, or generate one
            ssh_key = self._ssh_key
            if ssh_key:
                if not ssh_key.exists():
                    raise FileNotFoundError(f"{ssh_key} specified, but does not exist")
                ssh_pubkey = Path(f"{ssh_key}.pub")
            else:
                with self.lock:
                    self._tmpdir = tempfile.TemporaryDirectory()
                    ssh_key, ssh_pubkey = util.ssh_keygen(self._tmpdir.name)

            pubkey_contents = ssh_pubkey.read_text().strip()
            # TODO: split ^^^ into 3 parts (key type, hash, comment), assert it,
            #       and anonymize comment in case it contains a secret user/hostname
            spec_env["secrets"]["RESERVE_SSH_PUBKEY"] = pubkey_contents

            with self.lock:
                self.request = Request(api=self.api)
                self.request.submit(spec)
            logger.info(f"submitted request {self.request.id}")
            logger.debug(
                f"request {self.request.id}:\n{textwrap.indent(str(self.request), '    ')}",
            )

            # wait for user/host to ssh to
            ssh_user = ssh_host = None
            for line in PipelineLogStreamer(self.request):
                # the '\033[0m' is to reset colors sometimes left in a bad
                # state by pipeline.log
                logger.debug(f"{line}\033[0m")
                # find hidden login details
                m = re.search(
                    # host address can be an IP address or a hostname
                    r"\] Guest is ready: ArtemisGuest\([^,]+, (\w+)@([^,]+), arch=",
                    line,
                )
                if m:
                    ssh_user, ssh_host = m.groups()
                    continue
                # but wait until much later despite having login, at least until
                # the test starts running (and we get closer to it inserting our
                # ~/.ssh/authorized_keys entry)
                if ssh_user and re.search(r"\] starting tests execution", line):
                    break

            # wait for a successful connection over ssh
            # (it will be failing to login for a while, until the reserve test
            #  installs our ssh pubkey into authorized_keys)
            ssh_attempt_cmd = (
                "ssh", "-q", "-i", ssh_key.absolute(), "-oConnectionAttempts=60",
                "-oStrictHostKeyChecking=no", "-oUserKnownHostsFile=/dev/null",
                "-oBatchMode=yes",
                f"{ssh_user}@{ssh_host}", "exit 123",
            )
            while True:
                time.sleep(1)
                self.request.assert_alive()

                proc = util.subprocess_run(
                    ssh_attempt_cmd,
                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                )
                if proc.returncode == 123:
                    break

            return self.Reserved(
                host=ssh_host,
                port=22,
                user=ssh_user,
                ssh_key=ssh_key,
                request=self.request,
            )

        except:
            self.release()
            raise

    def release(self):
        with self.lock:
            if self.request:
                try:
                    self.request.cancel()
                except APIError:
                    pass
                finally:
                    self.request = None

            if self._tmpdir:
                self._tmpdir.cleanup()
                self._tmpdir = None

    def __enter__(self):
        try:
            return self.reserve()
        except Exception:
            self.release()
            raise

    def __exit__(self, exc_type, exc_value, traceback):
        self.release()
