#!/usr/bin/python3
"""CUPS filter for TA 3505ci - injects department code into PostScript stream."""
import sys


def parse_options(options_str):
    options = {}
    for token in options_str.split():
        if '=' in token:
            k, v = token.split('=', 1)
            options[k] = v
        else:
            options[token] = 'true'
    return options


def get_account_code(options):
    km = options.get('KmManagment', 'Default')
    if km == 'Default' or not km.startswith('MG'):
        return None
    code = km[2:]
    if not code.isdigit():
        return None
    return code  # raw digits, no zero-padding — matches kyofilter_E behaviour


def process_stream(lines, account_code):
    if not account_code:
        yield from lines
        return

    inject = f"({account_code}) statusdict /setmanagementnumber get exec\n".encode()
    injected = False

    for line in lines:
        stripped = line.rstrip(b'\r\n')
        if not injected and (stripped.rstrip() == b'%%EndSetup' or stripped.startswith(b'%%Page:')):
            yield inject
            injected = True
        yield line

    if not injected:
        sys.stderr.write("kyofilter: WARNING: no injection point found in PostScript stream\n")



def main():
    if len(sys.argv) < 6:
        sys.stderr.write("Usage: kyofilter job-id user title copies options [file]\n")
        sys.exit(1)

    options = parse_options(sys.argv[5])
    account_code = get_account_code(options)

    infile = open(sys.argv[6], 'rb') if len(sys.argv) > 6 else sys.stdin.buffer
    try:
        for chunk in process_stream(infile, account_code):
            sys.stdout.buffer.write(chunk)
    finally:
        if len(sys.argv) > 6:
            infile.close()


if __name__ == '__main__':
    main()
