#  Copyright 2024 Vsevolod Stakhov
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

from urllib.request import urlopen
import glob
import grp
import http.client
import os
import os.path
import psutil
import pwd
import shutil
import signal
import socket
import ssl
import stat
import subprocess
import random
import re
import sys
import tempfile

from robot.api import logger
from robot.libraries.BuiltIn import BuiltIn
import json


def Check_JSON(j):
    d = json.JSONDecoder(strict=True).decode(j.decode('utf-8'))
    logger.debug('got json %s' % d)
    assert len(d) > 0
    assert 'error' not in d
    return d


def check_json_log(fn):
    line_count = 0
    f = open(fn, 'r', encoding="utf-8")
    for l in f.readlines():
        d = json.JSONDecoder(strict=True).decode(l)
        assert len(d) > 0
        line_count = line_count + 1
    assert line_count > 0


def cleanup_temporary_directory(directory):
    shutil.rmtree(directory)


def save_run_results(directory, filenames):
    current_directory = os.getcwd()
    suite_name = BuiltIn().get_variable_value("${SUITE_NAME}")
    test_name = BuiltIn().get_variable_value("${TEST NAME}")
    if os.path.exists(directory):
        onlyfiles = [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))]
        logger.debug('%s content before cleanup: %s' % (directory, onlyfiles))
        if test_name is None:
            # this is suite-level tear down
            destination_directory = "%s/robot-save/%s" % (current_directory, suite_name)
        else:
            destination_directory = "%s/robot-save/%s/%s" % (current_directory, suite_name, test_name)
        if not os.path.isdir(destination_directory):
            os.makedirs(destination_directory)
        for file in filenames.split(' '):
            source_file = "%s/%s" % (directory, file)
            logger.debug('check if we can save %s' % source_file)
            if os.path.isfile(source_file):
                logger.debug('found %s, save it' % file)
                shutil.copy(source_file, "%s/%s" % (destination_directory, file))
                shutil.copy(source_file, "%s/robot-save/%s.last" % (current_directory, file))


def encode_filename(filename):
    return "".join(['%%%0X' % ord(b) for b in filename])


def get_test_directory():
    return os.path.abspath(os.path.dirname(os.path.realpath(__file__)) + "../../")


def get_top_dir():
    if os.environ.get('RSPAMD_TOPDIR'):
        return os.environ['RSPAMD_TOPDIR']

    return get_test_directory() + "/../../"


def get_install_root():
    if os.environ.get('RSPAMD_INSTALLROOT'):
        return os.path.abspath(os.environ['RSPAMD_INSTALLROOT'])

    return os.path.abspath("../install/")


def get_rspamd():
    if os.environ.get('RSPAMD'):
        return os.environ['RSPAMD']
    if os.environ.get('RSPAMD_INSTALLROOT'):
        return os.environ['RSPAMD_INSTALLROOT'] + "/bin/rspamd"
    dname = get_top_dir()
    return dname + "/src/rspamd"


def get_rspamc():
    if os.environ.get('RSPAMC'):
        return os.environ['RSPAMC']
    if os.environ.get('RSPAMD_INSTALLROOT'):
        return os.environ['RSPAMD_INSTALLROOT'] + "/bin/rspamc"
    dname = get_top_dir()
    return dname + "/src/client/rspamc"


def get_rspamadm():
    if os.environ.get('RSPAMADM'):
        return os.environ['RSPAMADM']
    if os.environ.get('RSPAMD_INSTALLROOT'):
        return os.environ['RSPAMD_INSTALLROOT'] + "/bin/rspamadm"
    dname = get_top_dir()
    return dname + "/src/rspamadm/rspamadm"


def HTTP(method, host, port, path, data=None, headers={}):
    c = http.client.HTTPConnection("%s:%s" % (host, port))
    c.request(method, path, data, headers)
    r = c.getresponse()
    t = r.read()
    s = r.status
    c.close()
    return [s, t]


def HTTP_With_Headers(method, host, port, path, data=None, headers={}):
    """HTTP request that returns response headers.
    Returns [status, body, headers_dict]
    """
    c = http.client.HTTPConnection("%s:%s" % (host, port))
    c.request(method, path, data, headers)
    r = c.getresponse()
    t = r.read()
    s = r.status
    h = dict(r.getheaders())
    c.close()
    return [s, t, h]


def HTTPS(method, host, port, path, data=None, headers={}):
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    c = http.client.HTTPSConnection("%s:%s" % (host, port), context=ctx)
    c.request(method, path, data, headers)
    r = c.getresponse()
    t = r.read()
    s = r.status
    c.close()
    return [s, t]


def HTTPS_With_Headers(method, host, port, path, data=None, headers={}):
    """HTTPS request that returns response headers.
    Returns [status, body, headers_dict]
    """
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    c = http.client.HTTPSConnection("%s:%s" % (host, port), context=ctx)
    c.request(method, path, data, headers)
    r = c.getresponse()
    t = r.read()
    s = r.status
    h = dict(r.getheaders())
    c.close()
    return [s, t, h]


def generate_ssl_cert(tmpdir):
    """Generate a self-signed EC certificate and key in tmpdir.
    Returns (cert_path, key_path).
    """
    cert_path = os.path.join(tmpdir, "test-cert.pem")
    key_path = os.path.join(tmpdir, "test-key.pem")
    subprocess.check_call([
        "openssl", "req", "-x509", "-newkey", "ec",
        "-pkeyopt", "ec_paramgen_curve:prime256v1",
        "-keyout", key_path, "-out", cert_path,
        "-days", "1", "-nodes",
        "-subj", "/CN=rspamd-test",
    ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    # Make readable by rspamd worker (runs as nobody)
    os.chmod(cert_path, 0o644)
    os.chmod(key_path, 0o644)
    return cert_path, key_path


def hard_link(src, dst):
    os.link(src, dst)


def make_temporary_directory():
    """Creates and returns a unique temporary directory

    Example:
    | ${RSPAMD_TMPDIR} = | Make Temporary Directory |
    """
    dirname = tempfile.mkdtemp()
    os.chmod(dirname, stat.S_IRUSR |
             stat.S_IXUSR |
             stat.S_IWUSR |
             stat.S_IRGRP |
             stat.S_IXGRP |
             stat.S_IROTH |
             stat.S_IXOTH)
    return dirname


def make_temporary_file():
    return tempfile.mktemp()


def path_splitter(path):
    dirname = os.path.dirname(path)
    basename = os.path.basename(path)
    return [dirname, basename]


def rspamc(addr, port, filename):
    mboxgoo = b"From MAILER-DAEMON Fri May 13 19:17:40 2016\r\n"
    goo = open(filename, 'rb').read()
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((addr, port))
    s.send(b"CHECK RSPAMC/1.0\r\nContent-length: ")
    s.send(str(len(goo + mboxgoo)).encode('utf-8'))
    s.send(b"\r\n\r\n")
    s.send(mboxgoo)
    s.send(goo)
    data = b""
    while True:
        chunk = s.recv(32768)
        if not chunk:
            break
        data += chunk
    s.close()
    return data.decode('utf-8')


def Scan_File(filename, **headers):
    addr = BuiltIn().get_variable_value("${RSPAMD_LOCAL_ADDR}")
    port = BuiltIn().get_variable_value("${RSPAMD_PORT_NORMAL}")
    headers["Queue-Id"] = BuiltIn().get_variable_value("${TEST_NAME}")
    c = http.client.HTTPConnection("%s:%s" % (addr, port))
    c.request("POST", "/checkv2", open(filename, "rb"), headers)
    r = c.getresponse()
    assert r.status == 200
    d = json.JSONDecoder(strict=True).decode(r.read().decode('utf-8'))
    c.close()
    BuiltIn().set_test_variable("${SCAN_RESULT}", d)
    return


def _build_multipart(boundary, metadata_json, message_bytes):
    """Build a multipart/form-data body with metadata and message parts."""
    body = b""
    body += ("--" + boundary + "\r\n").encode()
    body += b"Content-Disposition: form-data; name=\"metadata\"\r\n"
    body += b"Content-Type: application/json\r\n"
    body += b"\r\n"
    if isinstance(metadata_json, str):
        metadata_json = metadata_json.encode('utf-8')
    body += metadata_json
    body += b"\r\n"
    body += ("--" + boundary + "\r\n").encode()
    body += b"Content-Disposition: form-data; name=\"message\"\r\n"
    body += b"\r\n"
    if isinstance(message_bytes, str):
        message_bytes = message_bytes.encode('utf-8')
    body += message_bytes
    body += b"\r\n"
    body += ("--" + boundary + "--\r\n").encode()
    return body


def _build_multipart_single(boundary, part_name, part_data, content_type=None):
    """Build a multipart/form-data body with a single part."""
    body = b""
    body += ("--" + boundary + "\r\n").encode()
    body += ("Content-Disposition: form-data; name=\"%s\"\r\n" % part_name).encode()
    if content_type:
        body += ("Content-Type: %s\r\n" % content_type).encode()
    body += b"\r\n"
    if isinstance(part_data, str):
        part_data = part_data.encode('utf-8')
    body += part_data
    body += b"\r\n"
    body += ("--" + boundary + "--\r\n").encode()
    return body


def _parse_multipart_response(body, content_type):
    """Parse a multipart/mixed response and return the 'result' part data as string."""
    if isinstance(body, bytes):
        body = body.decode('utf-8', errors='replace')

    # Extract boundary from Content-Type header
    m = re.search(r'boundary="?([^";]+)"?', content_type)
    if not m:
        raise ValueError("No boundary found in Content-Type: %s" % content_type)
    boundary = m.group(1)

    # Split on boundary
    parts = body.split("--" + boundary)
    for part in parts:
        if part.startswith("--"):
            continue  # closing boundary
        if not part.strip():
            continue

        # Split headers from body
        if "\r\n\r\n" in part:
            headers, data = part.split("\r\n\r\n", 1)
        elif "\n\n" in part:
            headers, data = part.split("\n\n", 1)
        else:
            continue

        # Check if this is the "result" part
        if 'name="result"' in headers:
            # Strip trailing \r\n
            data = data.rstrip("\r\n")
            return data

    raise ValueError("No 'result' part found in multipart response")


def Scan_File_V3(filename, metadata=None, **headers):
    """Send a /checkv3 multipart request and set ${SCAN_RESULT}."""
    addr = BuiltIn().get_variable_value("${RSPAMD_LOCAL_ADDR}")
    port = BuiltIn().get_variable_value("${RSPAMD_PORT_NORMAL}")

    meta = metadata if metadata else {}
    meta_json = json.dumps(meta)
    message_data = open(filename, "rb").read()

    boundary = "----rspamd-test-%016x" % random.getrandbits(64)
    body = _build_multipart(boundary, meta_json, message_data)

    headers["Content-Type"] = "multipart/form-data; boundary=" + boundary
    if "Queue-Id" not in headers:
        headers["Queue-Id"] = BuiltIn().get_variable_value("${TEST_NAME}")

    c = http.client.HTTPConnection("%s:%s" % (addr, port))
    c.request("POST", "/checkv3", body, headers)
    r = c.getresponse()
    assert r.status == 200, "Expected HTTP 200 but got %d" % r.status

    resp_body = r.read()
    resp_ct = r.getheader("Content-Type", "")
    result_data = _parse_multipart_response(resp_body, resp_ct)

    d = json.JSONDecoder(strict=True).decode(result_data)
    c.close()
    BuiltIn().set_test_variable("${SCAN_RESULT}", d)
    return


def Scan_File_V3_Expect_Error(filename, expected_status, metadata=None,
                               body_override=None, content_type_override=None,
                               **headers):
    """Send a /checkv3 request and expect a specific HTTP error status."""
    addr = BuiltIn().get_variable_value("${RSPAMD_LOCAL_ADDR}")
    port = BuiltIn().get_variable_value("${RSPAMD_PORT_NORMAL}")

    boundary = "----rspamd-test-%016x" % random.getrandbits(64)

    if body_override is not None:
        body = body_override
    else:
        meta = metadata if metadata else {}
        meta_json = json.dumps(meta)
        message_data = open(filename, "rb").read() if filename else b""
        body = _build_multipart(boundary, meta_json, message_data)

    if content_type_override:
        headers["Content-Type"] = content_type_override
    else:
        headers["Content-Type"] = "multipart/form-data; boundary=" + boundary

    if "Queue-Id" not in headers:
        headers["Queue-Id"] = BuiltIn().get_variable_value("${TEST_NAME}")

    c = http.client.HTTPConnection("%s:%s" % (addr, port))
    c.request("POST", "/checkv3", body, headers)
    r = c.getresponse()
    actual_status = r.status
    r.read()
    c.close()
    assert actual_status == int(expected_status), \
        "Expected HTTP %s but got %d" % (expected_status, actual_status)
    return


def Scan_File_V3_Single_Part(part_name, part_data, content_type_part=None, **headers):
    """Send a /checkv3 request with only a single part."""
    addr = BuiltIn().get_variable_value("${RSPAMD_LOCAL_ADDR}")
    port = BuiltIn().get_variable_value("${RSPAMD_PORT_NORMAL}")

    boundary = "----rspamd-test-%016x" % random.getrandbits(64)
    body = _build_multipart_single(boundary, part_name, part_data, content_type_part)

    headers["Content-Type"] = "multipart/form-data; boundary=" + boundary
    if "Queue-Id" not in headers:
        headers["Queue-Id"] = BuiltIn().get_variable_value("${TEST_NAME}")

    c = http.client.HTTPConnection("%s:%s" % (addr, port))
    c.request("POST", "/checkv3", body, headers)
    r = c.getresponse()
    status = r.status
    r.read()
    c.close()
    return status


def Scan_File_SSL(filename, port=None, **headers):
    """Like Scan_File but over HTTPS (TLS) to the normal worker SSL port."""
    addr = BuiltIn().get_variable_value("${RSPAMD_LOCAL_ADDR}")
    if port is None:
        port = BuiltIn().get_variable_value("${RSPAMD_PORT_NORMAL_SSL}")
    headers["Queue-Id"] = BuiltIn().get_variable_value("${TEST_NAME}")
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    c = http.client.HTTPSConnection("%s:%s" % (addr, port), context=ctx)
    c.request("POST", "/checkv2", open(filename, "rb"), headers)
    r = c.getresponse()
    assert r.status == 200, "Expected HTTP 200 but got %d" % r.status
    d = json.JSONDecoder(strict=True).decode(r.read().decode('utf-8'))
    c.close()
    BuiltIn().set_test_variable("${SCAN_RESULT}", d)
    return


def Scan_File_V3_SSL(filename, port=None, metadata=None, **headers):
    """Like Scan_File_V3 but over HTTPS (TLS)."""
    addr = BuiltIn().get_variable_value("${RSPAMD_LOCAL_ADDR}")
    if port is None:
        port = BuiltIn().get_variable_value("${RSPAMD_PORT_NORMAL_SSL}")

    meta = metadata if metadata else {}
    meta_json = json.dumps(meta)
    message_data = open(filename, "rb").read()

    boundary = "----rspamd-test-%016x" % random.getrandbits(64)
    body = _build_multipart(boundary, meta_json, message_data)

    headers["Content-Type"] = "multipart/form-data; boundary=" + boundary
    if "Queue-Id" not in headers:
        headers["Queue-Id"] = BuiltIn().get_variable_value("${TEST_NAME}")

    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    c = http.client.HTTPSConnection("%s:%s" % (addr, port), context=ctx)
    c.request("POST", "/checkv3", body, headers)
    r = c.getresponse()
    assert r.status == 200, "Expected HTTP 200 but got %d" % r.status

    resp_body = r.read()
    resp_ct = r.getheader("Content-Type", "")
    result_data = _parse_multipart_response(resp_body, resp_ct)

    d = json.JSONDecoder(strict=True).decode(result_data)
    c.close()
    BuiltIn().set_test_variable("${SCAN_RESULT}", d)
    return


def Send_SIGUSR1(pid):
    pid = int(pid)
    os.kill(pid, signal.SIGUSR1)


def set_directory_ownership(path, username, groupname):
    if os.getuid() == 0:
        uid = pwd.getpwnam(username).pw_uid
        gid = grp.getgrnam(groupname).gr_gid
        os.chown(path, uid, gid)


def spamc(addr, port, filename):
    goo = open(filename, 'rb').read()
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((addr, port))
    s.send(b"SYMBOLS SPAMC/1.0\r\nContent-length: ")
    s.send(str(len(goo)).encode('utf-8'))
    s.send(b"\r\n\r\n")
    s.send(goo)
    s.shutdown(socket.SHUT_WR)
    data = b""
    while True:
        chunk = s.recv(32768)
        if not chunk:
            break
        data += chunk
    s.close()
    return data.decode('utf-8')


def TCP_Connect(addr, port):
    """Attempts to open a TCP connection to specified address:port

    Example:
    | Wait Until Keyword Succeeds | 5s | 10ms | TCP Connect | localhost | 8080 |
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(5)  # seconds
    s.connect((addr, port))
    s.close()


def port_is_free(addr, port):
    """Assert that nothing is listening on addr:port.

    Used by teardown to confirm a just-terminated rspamd has actually
    released its listening sockets before the next suite on this pabot
    worker reuses the same port range. `Wait For Process` only reaps the
    main rspamd; the listening sockets are shared with forked workers and
    can linger briefly after main exits. rspamd sets SO_REUSEADDR, so this
    is NOT about TIME_WAIT -- a still-LISTENing socket from a not-yet-gone
    worker genuinely fails the next bind() with EADDRINUSE. Connecting and
    succeeding means someone is still listening -> raise so Wait Until
    Keyword Succeeds retries; connection refused means the port is free.

    Example:
    | Wait Until Keyword Succeeds | 10s | 0.2s | Port Is Free | 127.0.0.1 | 25790 |
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(0.5)
    try:
        s.connect((addr, int(port)))
    except (ConnectionRefusedError, socket.timeout, OSError):
        return True
    finally:
        s.close()
    raise AssertionError("port %s:%s is still in use" % (addr, port))


def try_reap_zombies():
    try:
        os.waitpid(-1, os.WNOHANG)
    except ChildProcessError:
        pass


def ping_rspamd(addr, port):
    return str(urlopen("http://%s:%s/ping" % (addr, port)).read())


def redis_check(addr, port):
    """Attempts to open a TCP connection to specified address:port

    Example:
    | Wait Until Keyword Succeeds | 5s | 10ms | TCP Connect | localhost | 8080 |
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(1.0)  # seconds
    s.connect((addr, port))
    if s.sendall(b"ECHO TEST\n"):
        result = s.recv(128)
        return result == b'TEST\n'
    else:
        return False


def update_dictionary(a, b):
    a.update(b)
    return a


TERM_TIMEOUT = 10  # wait after sending a SIGTERM signal
KILL_WAIT = 20  # additional wait after sending a SIGKILL signal


def shutdown_process(process):
    # send SIGTERM
    process.terminate()

    try:
        process.wait(TERM_TIMEOUT)
        return
    except psutil.TimeoutExpired:
        logger.info("PID {} is not terminated in {} seconds, sending SIGKILL...".format(process.pid, TERM_TIMEOUT))
        try:
            # send SIGKILL
            process.kill()
        except psutil.NoSuchProcess:
            # process exited just before we tried to kill
            return

    try:
        process.wait(KILL_WAIT)
    except psutil.TimeoutExpired:
        raise RuntimeError("Failed to shutdown process {} ({})".format(process.pid, process.name()))


def shutdown_process_with_children(pid):
    pid = int(pid)
    try:
        process = psutil.Process(pid=pid)
    except psutil.NoSuchProcess:
        return
    children = process.children(recursive=True)
    shutdown_process(process)
    for child in children:
        try:
            child.kill()
        except psutil.NoSuchProcess:
            pass
    psutil.wait_procs(children, timeout=KILL_WAIT)


def write_to_stdin(process_handle, text):
    if not isinstance(text, bytes):
        text = bytes(text, 'utf-8')
    lib = BuiltIn().get_library_instance('Process')
    obj = lib.get_process_object()
    obj.stdin.write(text + b"\n")
    obj.stdin.flush()
    obj.stdin.close()
    out = obj.stdout.read(4096)
    return out.decode('utf-8')


def get_file_if_exists(file_path):
    if os.path.exists(file_path):
        with open(file_path, 'r') as myfile:
            return myfile.read()
    return None


def _merge_luacov_stats(statsfile, coverage):
    """
    Reads a coverage stats file written by luacov and merges coverage data to
    'coverage' dict: { src_file: hits_list }

    Format of the file defined in:
    https://github.com/keplerproject/luacov/blob/master/src/luacov/stats.lua
    """
    with open(statsfile, 'r') as fh:
        while True:
            # max_line:filename
            line = fh.readline().rstrip()
            if not line:
                break

            max_line, src_file = line.split(':')
            counts = [int(x) for x in fh.readline().split()]
            assert len(counts) == int(max_line)

            if src_file in coverage:
                # enlarge list if needed: lenght of list in different luacov.stats.out files may differ
                old_len = len(coverage[src_file])
                new_len = len(counts)
                if new_len > old_len:
                    coverage[src_file].extend([0] * (new_len - old_len))
                # sum execution counts for each line
                for l, exe_cnt in enumerate(counts):
                    coverage[src_file][l] += exe_cnt
            else:
                coverage[src_file] = counts


def _dump_luacov_stats(statsfile, coverage):
    """
    Saves data to the luacov stats file. Existing file is overwritted if exists.
    """
    src_files = sorted(coverage)

    with open(statsfile, 'w') as fh:
        for src in src_files:
            stats = " ".join(str(n) for n in coverage[src])
            fh.write("%s:%s\n%s\n" % (len(coverage[src]), src, stats))


# File used by luacov to collect coverage stats
LUA_STATSFILE = "luacov.stats.out"


def collect_lua_coverage():
    """
    Merges ${RSPAMD_TMPDIR}/*.luacov.stats.out into luacov.stats.out

    Example:
    | Collect Lua Coverage |
    """
    # decided not to do optional coverage so far
    # if not 'ENABLE_LUA_COVERAGE' in os.environ['HOME']:
    #    logger.info("ENABLE_LUA_COVERAGE is not present in env, will not collect Lua coverage")
    #    return

    tmp_dir = BuiltIn().get_variable_value("${RSPAMD_TMPDIR}")

    coverage = {}
    input_files = []

    for f in glob.iglob("%s/*.luacov.stats.out" % tmp_dir):
        _merge_luacov_stats(f, coverage)
        input_files.append(f)

    if input_files:
        if os.path.isfile(LUA_STATSFILE):
            _merge_luacov_stats(LUA_STATSFILE, coverage)
        _dump_luacov_stats(LUA_STATSFILE, coverage)
        logger.info("%s merged into %s" % (", ".join(input_files), LUA_STATSFILE))
    else:
        logger.info("no *.luacov.stats.out files found in %s" % tmp_dir)


def file_exists(file):
    return os.path.isfile(file)


def redis_stream_read_msgpack(host, port, stream_key):
    """Read latest entry from Redis stream and decode msgpack data.

    Returns decoded dict with metadata fields.

    Example:
    | ${data} = | Redis Stream Read Msgpack | ${RSPAMD_REDIS_ADDR} | ${RSPAMD_REDIS_PORT} | test:structured |
    """
    try:
        import redis
    except ImportError:
        raise Exception("redis module not installed - run: pip install redis")

    try:
        import msgpack
    except ImportError:
        raise Exception("msgpack module not installed - run: pip install msgpack")

    r = redis.Redis(host=host, port=int(port), decode_responses=False)

    # Read from stream
    entries = r.xrange(stream_key, count=1)
    if not entries:
        raise Exception(f"No data in stream {stream_key}")

    # Get the first entry's data field
    entry_id, fields = entries[0]
    if b'data' not in fields:
        raise Exception(f"No data field in stream entry, keys: {list(fields.keys())}")

    msgpack_data = fields[b'data']

    # Decode msgpack with raw=True to preserve bytes, then convert what we can
    decoded = msgpack.unpackb(msgpack_data, raw=True)

    # Convert bytes keys to strings for easier access
    def convert_keys(obj):
        if isinstance(obj, dict):
            return {k.decode('utf-8') if isinstance(k, bytes) else k: convert_keys(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [convert_keys(item) for item in obj]
        elif isinstance(obj, bytes):
            # Try to decode as UTF-8, otherwise keep as bytes
            try:
                return obj.decode('utf-8')
            except UnicodeDecodeError:
                return obj
        return obj

    return convert_keys(decoded)


def validate_structured_metadata(data, expected_fields=None):
    """Validate structured metadata export format.

    Checks that required fields exist and UUID v7 has correct format.

    Example:
    | Validate Structured Metadata | ${data} | uuid,ip,score,action |
    """
    import re

    if expected_fields is None:
        expected_fields = 'uuid,ip,score,action'

    errors = []

    for field in expected_fields.split(','):
        field = field.strip()
        if field not in data:
            errors.append(f"Missing field: {field}")

    # Validate UUID v7 format if present
    if 'uuid' in data:
        uuid = data['uuid']
        # UUID v7: xxxxxxxx-xxxx-7xxx-xxxx-xxxxxxxxxxxx
        if not re.match(r'^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$', str(uuid)):
            errors.append(f"Invalid UUID v7 format: {uuid}")

    if errors:
        raise Exception("Validation errors: " + "; ".join(errors))

    return True


def validate_zstd_compressed_fields(data):
    """Validate that zstd compression markers are set correctly.

    Returns count of compressed fields found.

    Example:
    | ${count} = | Validate Zstd Compressed Fields | ${data} |
    """
    count = 0

    # Check text_compressed flag
    if data.get('text_compressed'):
        count += 1

    # Check attachments
    for att in data.get('attachments', []):
        if att.get('content_compressed'):
            count += 1

    # Check images
    for img in data.get('images', []):
        if img.get('content_compressed'):
            count += 1

    return count


def validate_attachments_have_content_type(data):
    """Validate that attachments have content_type field.

    Returns count of attachments with content_type.

    Example:
    | ${count} = | Validate Attachments Have Content Type | ${data} |
    """
    count = 0
    for att in data.get('attachments', []):
        if 'content_type' in att and att['content_type']:
            count += 1
    return count
