import json
import logging
import os
import pwd
import re
import signal
import socket
import subprocess
import threading
import time
from dataclasses import dataclass
from typing import Optional, Dict, List

from clcommon.clpwd import drop_user_privileges

from clwpos import socket_utils
from clwpos.utils import run_in_cagefs_if_needed

_ADVICE_ID_RE = re.compile(r'^[1-9][0-9]*\Z')

_SMART_ADVICE_USER_UTILITY = '/opt/alt/php-xray/cl-smart-advice-user'

# Maps a php-xray advice `type` (as returned by `cl-smart-advice-user details`)
# to the AWP feature/module the advice belongs to. Mirrors php-xray's
# advice_mapping (xray/adviser/advice_types/__init__.py). Any advice type not
# listed here is treated as not resolvable and the apply is denied (fail-closed),
# so a new downstream advice type can only over-deny until this map is updated.
_ADVICE_TYPE_TO_FEATURE = {
    'OBJECT_CACHE': 'object_cache',
    'CDN': 'cdn',
    'CPCSS': 'critical_css',
    'IMAGE_OPTIMIZATION': 'image_optimization',
    # get_allowed_modules surfaces this entitlement as the Feature value
    # 'site_optimization' (SITE_OPTIMIZATION_FEATURE = Feature("site_optimization"),
    # features.py:1468), NOT the AWPSuite name 'accelerate_wp'.
    'SITE_OPTIMIZATION': 'site_optimization',
}


@dataclass
class EnableFeatureTask:
    uid: int

    # some arguments that we need when
    # we enable module manually
    domain: str
    wp_path: str

    # the features list that we would like to enable
    feature: list

    # if true, appends --ignore-errors to enable command
    ignore_errors: bool

    # time when we give up and
    # think that waiting is unsuccessful
    timeout_timestamp: int

    # optional argument which indicates that
    # we are applying advice
    advice_id: Optional[str] = None

    # pid of the process which took the task in work
    pid: Optional[int] = None

    # set under the handler lock once a handler thread has claimed this task and
    # is about to fork a child for it; prevents a second concurrent handler
    # thread from forking a duplicate child for the same still-pending task.
    in_progress: bool = False


class PendingSubscriptionWatcher:
    """
    Listen to set-suite and user events in order to automatically
    enable module when suite is purchased.
    """

    # max time that we wait for the feature to be allowed
    _FEATURE_STATUS_CHECK_TIMEOUT = 180

    _MAX_PENDING_TASKS_PER_UID = 10
    _MAX_WP_PATH_LENGTH = 256

    # how often the parent polls the forked child while waiting for it to exit
    _CHILD_POLL_INTERVAL = 0.5
    # grace given to a child after SIGTERM before escalating to SIGKILL
    _CHILD_TERM_GRACE = 5
    # per-CLI-invocation subprocess timeout enforced inside the child; the child
    # runs one bounded CLI call per feature sequentially.
    _CHILD_WAIT_TIMEOUT = 180
    # extra slack added on top of the child's worst-case total runtime so the
    # parent does not race the child's own per-invocation timeouts.
    _CHILD_WAIT_MARGIN = 30

    def __init__(self, logger=None):
        # list of pending tasks to automatically enable
        # module when suite is allowed
        self._pending_enable_tasks: Dict[int, List[EnableFeatureTask]] = dict()

        # The daemon spawns one handler thread per accepted connection and
        # permits several concurrent connections per uid, so every read/write of
        # _pending_enable_tasks (and the per-task claim) races. Guard them all
        # with one re-entrant lock (re-entrant because the cap-checking path
        # calls _cleanup_pending_tasks while already holding it). The lock is
        # only ever held around in-memory dict access - never across a
        # subprocess spawn or os.fork().
        self._tasks_lock = threading.RLock()

        self._logger = logger or logging.getLogger(__name__)

    @staticmethod
    def _known_features():
        """The server-defined set of valid feature interface names (~5 elements)."""
        from clwpos.optimization_features.features import ALL_OPTIMIZATION_FEATURES

        return {f.to_interface_name() for f in ALL_OPTIMIZATION_FEATURES}

    @staticmethod
    def _pid_alive(pid):
        """True if a child pid has been assigned and is still running."""
        if not pid:
            return False
        try:
            os.kill(pid, 0)
            return True
        except OSError:
            return False

    def _cleanup_pending_tasks(self):
        """
        Cleanup list of pending tasks by removing outdated.
        """
        with self._tasks_lock:
            for uid, pending_tasks in list(self._pending_enable_tasks.items()):
                for pending_task in pending_tasks:
                    if pending_task.timeout_timestamp >= time.time():
                        continue
                    # A live child must still be retained until it exits, even
                    # past the timeout.
                    if self._pid_alive(pending_task.pid):
                        continue
                    # A claimed task (in_progress) PAST its timeout with NO live
                    # child (pid is None or dead) is a stale claim - e.g. the
                    # parent died after the claim but before/around the fork,
                    # leaving no child to release it. Do not retain it forever:
                    # reset the stale claim and remove it per the existing
                    # timeout semantics so a later callback can re-queue/re-attempt
                    # rather than wedge the queue.
                    if pending_task.in_progress:
                        pending_task.in_progress = False
                        pending_task.pid = None
                    pending_tasks.remove(pending_task)
                if not pending_tasks:
                    self._pending_enable_tasks.pop(uid)
            # log only aggregate counts, not the cross-UID task table
            self._logger.info(
                'Cleanup of pending tasks. Still active: %d task(s) across %d uid(s)',
                sum(len(tasks) for tasks in self._pending_enable_tasks.values()),
                len(self._pending_enable_tasks),
            )

    def _run_and_log_results(self, args):
        self._logger.info('Running %s', ' '.join(args))
        try:
            # Bound the user CLI so an attacker-controlled wp_path (FIFO,
            # unresponsive path) cannot make it block forever in the child.
            output = run_in_cagefs_if_needed(args, check=True, timeout=self._CHILD_WAIT_TIMEOUT)
            self._logger.info('Command succeded with output: \n`%s`', output)
        except subprocess.CalledProcessError as e:
            self._logger.exception(
                'Unable to activate feature in background. Stdout is %s. Stderr is %s', e.stdout, e.stderr
            )
        except Exception:
            self._logger.exception('Unable to activate feature in background')
        self._logger.debug('Finished %s', ' '.join(args))

    def _resolve_advice_feature(self, advice_id):
        """
        Resolve advice_id -> the AWP feature it really belongs to via the
        read-only `cl-smart-advice-user details` query, run in the (already
        privilege-dropped) caller's context so php-xray's owner check applies.

        Returns the AWP feature name, or None if the advice cannot be resolved
        (owner mismatch / unknown id / unparsable response / unknown type).
        The caller MUST treat None as "deny" - fail closed.
        """
        try:
            output = run_in_cagefs_if_needed(
                [_SMART_ADVICE_USER_UTILITY, 'details', '--advice_id', str(advice_id)],
                check=True,
                timeout=self._CHILD_WAIT_TIMEOUT,
            ).stdout
            resp = json.loads(output)
        except subprocess.TimeoutExpired:
            self._logger.warning('advice details query timed out for advice_id=%s; refusing to apply', advice_id)
            return None
        except subprocess.CalledProcessError as e:
            self._logger.warning(
                'advice details query failed for advice_id=%s. Stdout is %s. Stderr is %s',
                advice_id,
                e.stdout,
                e.stderr,
            )
            return None
        except (ValueError, TypeError):
            self._logger.warning('advice details query returned unparsable output for advice_id=%s', advice_id)
            return None

        if not isinstance(resp, dict) or resp.get('result') != 'success':
            self._logger.warning('advice details query was not successful for advice_id=%s: %s', advice_id, resp)
            return None
        try:
            advice_type = resp['data']['advice']['type']
        except (KeyError, TypeError):
            self._logger.warning('advice details response missing advice type for advice_id=%s', advice_id)
            return None

        feature = _ADVICE_TYPE_TO_FEATURE.get(advice_type)
        if feature is None:
            self._logger.warning('advice_id=%s has unrecognised type %s; refusing to apply', advice_id, advice_type)
        return feature

    def _run_enable_command(self, task, allowed_features):
        if task.advice_id:
            # Bind the apply to the advice's TRUE feature, not the caller-declared
            # one: in advice mode only --advice_id is forwarded downstream, and the
            # advice's real feature is independent of task.feature - so a free
            # declared feature must not authorise applying a premium advice the
            # caller is not entitled to. Resolve the real feature and require it
            # to be in the caller's allowed set; deny (fail closed) otherwise.
            true_feature = self._resolve_advice_feature(task.advice_id)
            if true_feature is None or true_feature not in allowed_features:
                self._logger.info(
                    'not applying advice in background: advice_id=%s resolved feature=%s is not allowed for the caller',
                    task.advice_id,
                    true_feature,
                )
                return
            self._run_and_log_results([_SMART_ADVICE_USER_UTILITY, 'apply', '--advice_id', str(task.advice_id)])
        elif task.domain:
            for f in task.feature:
                self._logger.info('activate feature in background: feature=%s', f)
                cmd = [
                    'cloudlinux-awp-user',
                    'enable',
                    '--feature',
                    f,
                    '--domain',
                    task.domain,
                    '--wp-path',
                    task.wp_path,
                ]
                if task.ignore_errors:
                    cmd.append('--ignore-errors')
                self._run_and_log_results(cmd)
        else:
            logging.info('Task does not have any advice or domain specified, skipping')

    def _child_runtime_budget(self, task):
        """
        Worst-case wall-clock the forked child may legitimately need.

        The child (``_run_enable_command``) runs one bounded CLI invocation per
        feature sequentially, each capped at ``_CHILD_WAIT_TIMEOUT`` by its own
        subprocess timeout. The advice path runs TWO sequential bounded calls -
        the read-only ``cl-smart-advice-user details`` resolve followed by the
        ``apply`` - so it is budgeted for two per-CLI windows, not one.

        For the plain enable path the number of invocations is derived from the
        count of DISTINCT KNOWN features, clamped to the size of the known-feature
        set. task.feature is attacker-supplied; although the IPC path now dedups
        and caps it, deriving the budget from |distinct known features| keeps it
        capped by a small server constant (_CHILD_WAIT_TIMEOUT * |known_features|
        + margin) regardless of input.
        """
        if task.advice_id:
            # details resolve + apply, each capped at _CHILD_WAIT_TIMEOUT
            invocations = 2
        elif not task.feature:
            invocations = 1
        else:
            known_features = self._known_features()
            distinct_known = len(set(task.feature) & known_features)
            # clamp: never exceed the size of the known-feature set
            invocations = max(1, min(distinct_known, len(known_features)))
        return self._CHILD_WAIT_TIMEOUT * invocations + self._CHILD_WAIT_MARGIN

    def _reap_child(self, pid):
        """Reap a child pid that is known to have exited, ignoring ECHILD."""
        try:
            os.waitpid(pid, 0)
        except OSError:
            pass

    def _wait_for_child(self, pid, deadline):
        """
        Wait for a forked child to exit, bounded by ``deadline`` (epoch).

        Polls with WNOHANG instead of blocking on os.waitpid(pid, 0): a child
        running the user CLI against an attacker-controlled wp_path (e.g. a FIFO
        or an unresponsive path) could otherwise hang forever and stall the
        sequential pending-task loop for every other user. On deadline the child
        is SIGTERM'd, given a short grace, then SIGKILL'd, and reaped to avoid a
        zombie.
        """
        while True:
            try:
                reaped, _ = os.waitpid(pid, os.WNOHANG)
            except OSError:
                # child already gone (e.g. ECHILD); nothing to reap
                return
            if reaped:
                return
            if time.time() >= deadline:
                break
            time.sleep(self._CHILD_POLL_INTERVAL)

        self._logger.warning('background process pid=%d exceeded its deadline; terminating', pid)
        try:
            os.kill(pid, signal.SIGTERM)
        except OSError:
            self._reap_child(pid)
            return

        kill_deadline = time.time() + self._CHILD_TERM_GRACE
        while time.time() < kill_deadline:
            try:
                reaped, _ = os.waitpid(pid, os.WNOHANG)
            except OSError:
                return
            if reaped:
                return
            time.sleep(self._CHILD_POLL_INTERVAL)

        self._logger.warning('background process pid=%d did not exit after SIGTERM; killing', pid)
        try:
            os.kill(pid, signal.SIGKILL)
        except OSError:
            pass
        self._reap_child(pid)

    def suite_allowed_callback(self, client_socket_obj: socket.socket, daemon_socket: socket.socket = None):
        self._cleanup_pending_tasks()

        # Snapshot the (uid, tasks) pairs under the lock so the dict read is
        # synchronized, then fork/wait OUTSIDE the lock - the lock must never be
        # held across os.fork(), or the forked child would inherit a locked
        # mutex it can never release. The list mutations below are taken back
        # under the lock.
        with self._tasks_lock:
            snapshot = list(self._pending_enable_tasks.items())

        for uid, tasks in snapshot:
            # Iterate over a copy of the task list so each completed task can be
            # removed without mutating the list the for-loop walks. Whole-uid
            # del was unsafe when multiple tasks are pending for a single uid
            # (bugbot 9e736ea9 on MR !39).
            for task in list(tasks):
                # Atomically claim the task before forking: a second concurrent
                # handler thread must not fork a duplicate child for a task whose
                # first child is still running (the task stays pending across the
                # whole wait, and _cleanup_pending_tasks deliberately keeps a
                # live/in-progress task). Skip if already claimed; otherwise mark
                # it in_progress and proceed. The lock is NOT held across the fork
                # / wait below - a child must not inherit a held lock, and other
                # UIDs must not block for the whole wait (bugbot 412f1e5b on MR !60).
                with self._tasks_lock:
                    if task.in_progress or self._pid_alive(task.pid):
                        continue
                    task.in_progress = True
                # If os.fork() itself fails, no child exists: release the claim
                # and leave the task pending so a later callback can re-attempt
                # it. Otherwise the task would stay claimed (in_progress=True)
                # with no live child, _cleanup_pending_tasks would keep it, and
                # every later callback would skip it on the claim guard - a
                # permanent wedge.
                try:
                    fp = os.fork()
                except OSError:
                    self._logger.exception(
                        'fork failed for pending task uid=%s; leaving it pending for retry', task.uid
                    )
                    with self._tasks_lock:
                        task.pid = None
                        task.in_progress = False
                    continue
                if fp:
                    # Parent. Guarantee the claim is released even if the bounded
                    # wait raises, so the task can never wedge with the claim set
                    # and no live child.
                    try:
                        self._logger.info('background process forked: pid=%d', fp)
                        task.pid = fp
                        # Bound the wait so a single hung child (user-controlled
                        # wp_path) cannot freeze the loop for all other users, while
                        # still allowing a legitimate multi-feature enable its full
                        # per-feature budget (the child runs one bounded CLI call per
                        # feature sequentially). The bound is server-derived, so it
                        # stays finite regardless of attacker input.
                        deadline = time.time() + self._child_runtime_budget(task)
                        self._wait_for_child(fp, deadline)
                        self._logger.info('background process finished: pid=%d', fp)
                    finally:
                        with self._tasks_lock:
                            task.in_progress = False
                            # Another handler thread's _cleanup_pending_tasks may
                            # have already dropped this (now-finished) task, so
                            # removing it again would raise ValueError. Make the
                            # removal idempotent.
                            if task in tasks:
                                tasks.remove(task)
                else:
                    try:
                        if daemon_socket is not None:
                            daemon_socket.close()
                        client_socket_obj.close()
                        from clwpos.feature_suites import get_allowed_modules

                        # drop privileges forever.. or at least till the end of process
                        drop_user_privileges(pwd.getpwuid(task.uid).pw_name, effective_or_real=True, set_env=True)
                        allowed_features = get_allowed_modules(task.uid)

                        disallowed = [f for f in task.feature if f not in allowed_features]
                        for f in disallowed:
                            self._logger.info('unable to activate feature in background: feature=%s is not allowed', f)
                        task.feature = [f for f in task.feature if f in allowed_features]
                        # In advice mode the apply is gated on the advice's TRUE
                        # feature inside _run_enable_command, not on task.feature;
                        # the declared-feature filter below only gates the plain
                        # enable path.
                        if not task.advice_id and not task.feature:
                            os._exit(0)
                            return  # unreachable in production; keeps control flow explicit for test mocks

                        self._run_enable_command(task, allowed_features)
                    finally:
                        os._exit(0)
            with self._tasks_lock:
                if uid in self._pending_enable_tasks and not self._pending_enable_tasks[uid]:
                    del self._pending_enable_tasks[uid]

        response: dict = {"result": "success"}
        socket_utils.send_dict_to_socket_connection_and_close(client_socket_obj, response)

    def add_pending_upgrade_task(self, client_socket_obj: socket.socket, user_request: dict, uid: int):
        from clcommon.cpapi import get_main_username_by_uid, userdomains

        known_features = self._known_features()
        # Bound the raw feature string before splitting: a non-root peer can
        # otherwise send "cdn,cdn,..." that .split(',') explodes into a huge
        # list of valid-but-duplicate names, exhausting the shared root daemon's
        # memory. The longest legitimate request lists every known feature once,
        # so cap the wire size to that (names + comma separators).
        max_feature_string_length = len(known_features) * (max((len(f) for f in known_features), default=0) + 1)
        if len(user_request['feature']) > max_feature_string_length:
            response: dict = {"result": "error", "context": "feature list is too long"}
            socket_utils.send_dict_to_socket_connection_and_close(client_socket_obj, response)
            return
        # Deduplicate while preserving order: the stored list can never exceed
        # the number of known features, so duplicate names cannot inflate it (and
        # cannot inflate the parent's per-task child-runtime budget either).
        requested_features = list(dict.fromkeys(user_request['feature'].split(',')))
        for f in requested_features:
            if f not in known_features:
                response: dict = {"result": "error", "context": "unknown feature: %s" % f}
                socket_utils.send_dict_to_socket_connection_and_close(client_socket_obj, response)
                return

        advice_id = user_request.get('advice_id')
        if advice_id is not None:
            advice_id = str(advice_id)
            if not _ADVICE_ID_RE.match(advice_id):
                response: dict = {"result": "error", "context": "invalid advice_id format"}
                socket_utils.send_dict_to_socket_connection_and_close(client_socket_obj, response)
                return
            # Entitlement for advice mode is enforced at apply time against the
            # advice's TRUE feature (resolved via `cl-smart-advice-user details`
            # in _run_enable_command), not against the caller-declared feature -
            # the two are independent, so the declared feature cannot be relied on
            # to authorise the apply.

        wp_path = user_request['wp_path']
        if wp_path is not None and len(wp_path) > self._MAX_WP_PATH_LENGTH:
            response: dict = {"result": "error", "context": "wp_path is too long"}
            socket_utils.send_dict_to_socket_connection_and_close(client_socket_obj, response)
            return
        if wp_path and '..' in wp_path.split('/'):
            response: dict = {"result": "error", "context": "invalid wp_path: must not contain '..' components"}
            socket_utils.send_dict_to_socket_connection_and_close(client_socket_obj, response)
            return

        domain = user_request['domain']
        # Mirror the CLI's domain/wp_path coupling (wpos_user.py:444-447): both
        # must be set together. Without this, an IPC client that bypasses the
        # CLI can queue a task with truthy domain + wp_path=None, which crashes
        # _run_enable_command on `' '.join([..., '--wp-path', None])` and silently
        # fails feature activation.
        if bool(domain) != (wp_path is not None):
            response: dict = {"result": "error", "context": "domain and wp_path must be set together"}
            socket_utils.send_dict_to_socket_connection_and_close(client_socket_obj, response)
            return
        if domain:
            username = get_main_username_by_uid(uid)
            user_domains = {d for d, _ in userdomains(username)}
            if domain not in user_domains:
                response: dict = {"result": "error", "context": "domain does not belong to the user"}
                socket_utils.send_dict_to_socket_connection_and_close(client_socket_obj, response)
                return

        # Cleanup, dedup check, cap check and append must be one atomic critical
        # section: with one handler thread per connection (up to several per uid)
        # racing here, an unsynchronized check-then-append lets concurrent threads
        # each observe len(pending) < cap and all append, bypassing the per-uid
        # cap. Decide the outcome under the lock, then send the reply outside it
        # so no socket I/O happens while the lock is held.
        with self._tasks_lock:
            self._cleanup_pending_tasks()

            if uid not in self._pending_enable_tasks:
                self._pending_enable_tasks[uid] = []

            pending = self._pending_enable_tasks[uid]

            duplicate = any(
                existing.domain == domain
                and existing.wp_path == wp_path
                and set(existing.feature) == set(requested_features)
                and existing.advice_id == advice_id
                and existing.ignore_errors == user_request['ignore_errors']
                for existing in pending
            )

            if duplicate:
                response = {"result": "success"}
            elif len(pending) >= self._MAX_PENDING_TASKS_PER_UID:
                response = {"result": "error", "context": "too many pending tasks"}
            else:
                pending.append(
                    EnableFeatureTask(
                        uid=uid,
                        domain=domain,
                        wp_path=wp_path,
                        feature=requested_features,
                        ignore_errors=user_request['ignore_errors'],
                        advice_id=advice_id,
                        timeout_timestamp=int(time.time() + self._FEATURE_STATUS_CHECK_TIMEOUT),
                    )
                )
                self._logger.info('Successfully added pending upgrade subscription task %s', pending)
                response = {"result": "success"}

        socket_utils.send_dict_to_socket_connection_and_close(client_socket_obj, response)

    def get_upgrade_task_status(self, client_socket_obj: socket.socket, uid: int, feature: str):
        with self._tasks_lock:
            self._cleanup_pending_tasks()
            pending_tasks = self._pending_enable_tasks.get(uid, [])
            pending = any(feature in pending_task.feature for pending_task in pending_tasks)

        response: dict = {"result": "success", "pending": bool(pending)}
        socket_utils.send_dict_to_socket_connection_and_close(client_socket_obj, response)