from __future__ import annotations
import base64
import binascii
import io
import os
import sys
import zipfile
SOLUTION_OUTPUT_SEPARATOR: bytes = b'\n----- ENIGMA-SOLUTION-OUTPUT-BEGIN-a8c7f3e2-9d4b-4c5a-8f1e-2b6d3a4e5f7c -----\n'

def build_solution_zip(files: dict[str, str | bytes]) -> bytes:
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, 'w', compression=zipfile.ZIP_DEFLATED) as zf:
        for name, content in files.items():
            zf.writestr(name, content)
    return buffer.getvalue()

def write_solution_output(zip_bytes: bytes) -> None:
    sys.stdout.flush()
    sys.stderr.flush()
    encoded = base64.b64encode(zip_bytes)
    buf = sys.stdout.buffer
    buf.write(SOLUTION_OUTPUT_SEPARATOR)
    buf.write(encoded)
    buf.write(b'\n')
    buf.flush()

def split_on_separator(raw_stdout: bytes) -> tuple[bytes, bytes, bool]:
    idx = raw_stdout.find(SOLUTION_OUTPUT_SEPARATOR)
    if idx == -1:
        return (raw_stdout, b'', False)
    return (raw_stdout[:idx], raw_stdout[idx + len(SOLUTION_OUTPUT_SEPARATOR):], True)

def extract_artifacts(raw_stdout: bytes, output_dir: str) -> tuple[bool, str | None]:
    os.makedirs(output_dir, exist_ok=True)
    logs_bytes, payload_b64, found = split_on_separator(raw_stdout)
    log_path = os.path.join(output_dir, 'stdout.log')
    with open(log_path, 'wb') as f:
        f.write(logs_bytes)
    if not found:
        return (False, 'No solution output separator found in container stdout. The solver must print logs, then the separator line, then a base64-encoded zip of result.json and other artifacts.')
    payload_b64 = payload_b64.strip()
    if not payload_b64:
        return (False, 'Separator found but base64 payload is empty.')
    try:
        payload_bytes = base64.b64decode(payload_b64, validate=True)
    except (binascii.Error, ValueError) as e:
        return (False, f'Base64 decode of solution payload failed: {e}')
    if not payload_bytes:
        return (False, 'Decoded payload is empty.')
    zip_path = os.path.join(output_dir, 'solution_artifacts.zip')
    with open(zip_path, 'wb') as f:
        f.write(payload_bytes)
    if not zipfile.is_zipfile(zip_path):
        return (False, 'Decoded payload is not a valid zip file.')
    try:
        with zipfile.ZipFile(zip_path, 'r') as zf:
            dest = os.path.abspath(output_dir)
            for member in zf.infolist():
                member_path = os.path.normpath(os.path.join(dest, member.filename))
                if not (member_path == dest or member_path.startswith(dest + os.sep)):
                    return (False, f"Zip member '{member.filename}' escapes output directory.")
            zf.extractall(dest)
    except zipfile.BadZipFile as e:
        return (False, f'Bad zip file: {e}')
    except Exception as e:
        return (False, f'Zip extraction failed: {e}')
    return (True, None)