import logging
import os
import platform
import subprocess
import threading
from contextlib import contextmanager
from time import sleep
import sys

import pytest

import requests
import yaml
from retrying import retry

from constants import METRICS_PORT, MAX_RETRIES, BACKOFF_SECS, GRACEFUL_SHUTDOWN_TIMEOUT, READER_THREAD_JOIN_TIMEOUT

class CloudflaredProcess:
    """
    Wrapper around a Popen process that continuously drains stdout and stderr
    in background threads to prevent OS pipe buffers from filling up and
    blocking the child process. Captured output is logged when the process
    is cleaned up.
    """

    def __init__(self, cmd, allow_input, capture_output):
        output = subprocess.PIPE if capture_output else subprocess.DEVNULL
        stdin = subprocess.PIPE if allow_input else None
        self.process = subprocess.Popen(cmd, stdin=stdin, stdout=output, stderr=subprocess.STDOUT)

        self._capture_output = capture_output
        self._stdout_lines = []
        self._threads = []
        if capture_output:
            self._threads.append(self._start_reader(self.process.stdout, self._stdout_lines))

    @staticmethod
    def _start_reader(pipe, sink):
        def _drain():
            for line in pipe:
                sink.append(line)
            pipe.close()
        t = threading.Thread(target=_drain, daemon=True)
        t.start()
        return t

    def terminate(self):
        """Terminate the process if it is still running."""
        if self.process.poll() is None:
            self.process.terminate()

    def cleanup(self):
        """Terminate, wait for exit, join reader threads, and log output."""
        self.terminate()
        try:
            self.process.wait(timeout=GRACEFUL_SHUTDOWN_TIMEOUT)
        except subprocess.TimeoutExpired:
            self.process.kill()
            self.process.wait()
        for t in self._threads:
            t.join(timeout=READER_THREAD_JOIN_TIMEOUT)
        if self._capture_output:
            stdout = b"".join(self._stdout_lines).decode("utf-8", errors="replace")
            if stdout:
                LOGGER.info(f"cloudflared stdout:\n{stdout}")

    @property
    def stdout_lines(self):
        return self._stdout_lines

    # Proxy common Popen attributes so callers can still use the wrapper
    # as if it were a Popen (e.g. send_signal, stdin, pid, returncode).
    def __getattr__(self, name):
        return getattr(self.process, name)

def configure_logger():
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.DEBUG)
    handler = logging.StreamHandler(sys.stdout)
    logger.addHandler(handler)
    return logger

LOGGER = configure_logger()

def select_platform(plat):
    return pytest.mark.skipif(
        platform.system() != plat, reason=f"Only runs on {plat}")

def fips_enabled():
    env_fips = os.getenv("COMPONENT_TESTS_FIPS")
    return env_fips is not None and env_fips != "0"

nofips = pytest.mark.skipif(
        fips_enabled(), reason=f"Only runs without FIPS (COMPONENT_TESTS_FIPS=0)")

def skip_on_ci(reason):
    env_ci = os.getenv("CI")
    running_in_ci = env_ci is not None and env_ci != "0"
    return pytest.mark.skipif(
        running_in_ci, reason=f"This test can't run on CI due to: {reason}")

def write_config(directory, config):
    config_path = directory / "config.yml"
    with open(config_path, 'w') as outfile:
        yaml.dump(config, outfile)
    return config_path


def start_cloudflared(directory, config, cfd_args=["run"], cfd_pre_args=["tunnel"], new_process=False,
                      allow_input=False, capture_output=True, root=False, skip_config_flag=False, expect_success=True):

    config_path = None
    if not skip_config_flag:
        config_path = write_config(directory, config.full_config)

    cmd = cloudflared_cmd(config, config_path, cfd_args, cfd_pre_args, root)

    if new_process:
        return run_cloudflared_background(cmd, allow_input, capture_output)
    # By setting check=True, it will raise an exception if the process exits with non-zero exit code
    return subprocess.run(cmd, check=expect_success, capture_output=capture_output)

def cloudflared_cmd(config, config_path, cfd_args, cfd_pre_args, root):
    cmd = []
    if root:
        cmd += ["sudo"]
    cmd += [config.cloudflared_binary]
    cmd += cfd_pre_args

    if config_path is not None:
        cmd += ["--config", str(config_path)]

    cmd += cfd_args
    LOGGER.info(f"Run cmd {cmd} with config {config}")
    return cmd

@contextmanager
def run_cloudflared_background(cmd, allow_input, capture_output):
    cfd = None
    try:
        cfd = CloudflaredProcess(cmd, allow_input, capture_output)
        yield cfd
    finally:
        if cfd:
            cfd.cleanup()
    

def get_quicktunnel_url():
    quicktunnel_url = f'http://localhost:{METRICS_PORT}/quicktunnel'
    with requests.Session() as s:
        resp = send_request(s, quicktunnel_url, True)

        hostname = resp.json()["hostname"]
        assert hostname, \
            f"Quicktunnel endpoint returned {hostname} but we expected a url"

        return f"https://{hostname}"

def wait_tunnel_ready(tunnel_url=None, require_min_connections=1, cfd_logs=None):
    try:
        inner_wait_tunnel_ready(tunnel_url, require_min_connections)
    except Exception as e:
        if cfd_logs is not None:
            _log_cloudflared_logs(cfd_logs)
        raise e


@retry(stop_max_attempt_number=MAX_RETRIES, wait_fixed=BACKOFF_SECS * 1000)
def inner_wait_tunnel_ready(tunnel_url=None, require_min_connections=1):
    metrics_url = f'http://localhost:{METRICS_PORT}/ready'

    with requests.Session() as s:
        LOGGER.debug("Waiting for tunnel to be ready...")
        resp = send_request(s, metrics_url, True)

        ready_connections = resp.json()["readyConnections"]

        assert ready_connections >= require_min_connections, \
            f"Ready endpoint returned {resp.json()} but we expect at least {require_min_connections} connections"

        if tunnel_url is not None:
            send_request(s, tunnel_url, True)

def _log_cloudflared_logs(cfd_logs):
    log_file = cfd_logs
    if os.path.isdir(cfd_logs):
        files = os.listdir(cfd_logs)
        if len(files) == 0:
            return
        log_file = os.path.join(cfd_logs, files[0])
    with open(log_file, "r") as f:
        LOGGER.warning("Cloudflared Tunnel was not ready:")
        for line in f.readlines():
            LOGGER.warning(line)


@retry(stop_max_attempt_number=MAX_RETRIES, wait_fixed=BACKOFF_SECS * 1000)
def check_tunnel_not_connected():
    url = f'http://localhost:{METRICS_PORT}/ready'

    try:
        resp = requests.get(url, timeout=BACKOFF_SECS)
        assert resp.status_code == 503, f"Expect {url} returns 503, got {resp.status_code}"
        assert resp.json()[
            "readyConnections"] == 0, "Expected all connections to be terminated (pending reconnect)"
    # cloudflared might already terminate
    except requests.exceptions.ConnectionError as e:
        LOGGER.warning(f"Failed to connect to {url}, error: {e}")


def get_tunnel_connector_id():
    url = f'http://localhost:{METRICS_PORT}/ready'

    try:
        resp = requests.get(url, timeout=1)
        return resp.json()["connectorId"]
    # cloudflared might already terminated
    except requests.exceptions.ConnectionError as e:
        LOGGER.warning(f"Failed to connect to {url}, error: {e}")


# In some cases we don't need to check response status, such as when sending batch requests to generate logs
def send_requests(url, count, require_ok=True):
    errors = 0
    with requests.Session() as s:
        for _ in range(count):
            resp = send_request(s, url, require_ok)
            if resp is None:
                errors += 1
            sleep(0.01)
    if errors > 0:
        LOGGER.warning(
            f"{errors} out of {count} requests to {url} return non-200 status")


@retry(stop_max_attempt_number=MAX_RETRIES, wait_fixed=BACKOFF_SECS * 1000)
def send_request(session, url, require_ok):
    resp = session.get(url, timeout=BACKOFF_SECS)
    if require_ok:
        assert resp.status_code == 200, f"{url} returned {resp}"
    return resp if resp.status_code == 200 else None


def decode_jwt_payload(token):
    """
    Decode the payload section of a JWT token without signature verification.
    
    JWT Structure:
    ==============
    A JWT consists of three Base64URL-encoded parts separated by dots:
        HEADER.PAYLOAD.SIGNATURE
    
    The payload contains the JWT claims (the actual data/permissions).
    
    Args:
        token (str): The complete JWT token string
        
    Returns:
        dict: The decoded payload as a dictionary containing JWT claims
        
    Raises:
        ValueError: If the token doesn't have exactly 3 parts
        
    Note:
        This function does NOT verify the signature - it only decodes the payload.
        Use this only when you trust the token source (e.g., tokens you just generated).
    """
    import base64
    import json
    
    # Split JWT into its three components
    parts = token.split('.')
    if len(parts) != 3:
        raise ValueError(f"Invalid JWT format: expected 3 parts, got {len(parts)}")
    
    # Extract and decode the payload (middle section)
    # Base64 requires padding to be a multiple of 4 characters
    payload_encoded = parts[1]
    remainder = len(payload_encoded) % 4
    if remainder != 0:
        payload_padded = payload_encoded + '=' * (4 - remainder)
    else:
        payload_padded = payload_encoded
    
    # Decode from Base64URL format and parse JSON
    decoded_payload = base64.urlsafe_b64decode(payload_padded)
    return json.loads(decoded_payload)
