#!/usr/bin/env python3
"""BGClear CLI: Python 3.10+, no dependencies. Run with --help for instructions."""
import argparse
import email.utils
import glob
import http.client
import json
import math
import os
from pathlib import Path
import struct
import sys
import tempfile
import time
import urllib.error
import urllib.request
import uuid
import zlib

VERSION = '1.0.0'
API_URL = 'https://www.bgclear.ai/api/v1/remove'
MAX_INPUT = 25 * 1024 * 1024
MAX_OUTPUT = 48 * 1024 * 1024
EXTENSIONS = {'.png', '.jpg', '.jpeg', '.webp'}


class CliError(Exception):
    def __init__(self, message, fatal=False):
        super().__init__(message)
        self.fatal = fatal


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def parser():
    p = argparse.ArgumentParser(description='BGClear CLI — remove backgrounds and save transparent PNGs. Requires BGCLEAR_API_KEY. Free previews by default; images are uploaded for processing with store=false.')
    p.add_argument('inputs', nargs='+', help='Images, folders, or quoted glob patterns (PNG, JPEG, WebP)')
    p.add_argument('-o', '--output', type=Path, help='Exact .png output path for one image')
    p.add_argument('--output-dir', type=Path, default=Path('bgclear-output'), help='Output folder (default: bgclear-output)')
    p.add_argument('-r', '--recursive', action='store_true', help='Include subfolders and preserve their paths')
    p.add_argument('--overwrite', action='store_true', help='Replace existing outputs; never replace an input')
    p.add_argument('--size', choices=('preview', 'full'), default='preview', help='preview: free, longest side up to 800 px. full: 1 credit per successful image, up to 4 MP')
    p.add_argument('--version', action='version', version=f'BGClear CLI {VERSION}')
    return p


def plan(args):
    output_root = args.output_dir.resolve()
    found = {}
    for pattern in args.inputs:
        matches = [Path(pattern)] if Path(pattern).exists() else [Path(x) for x in sorted(glob.glob(pattern, recursive=args.recursive))]
        if not matches:
            raise CliError(f'No files match: {pattern}')
        for match in matches:
            if match.is_dir():
                paths = sorted(match.rglob('*') if args.recursive else match.iterdir())
                for path in paths:
                    if not path.is_file() or path.suffix.lower() not in EXTENSIONS or path.is_symlink():
                        continue
                    if output_root == path.resolve() or output_root in path.resolve().parents:
                        continue
                    found.setdefault(path.resolve(), path.relative_to(match))
            elif match.is_file():
                if match.suffix.lower() not in EXTENSIONS:
                    raise CliError(f'Unsupported input: {match}. Use PNG, JPEG or WebP.')
                found.setdefault(match.resolve(), Path(match.name))
    if not found:
        raise CliError('No supported images found. Use PNG, JPEG or WebP; check your output folder is separate from your input folder.')
    if args.output and len(found) != 1:
        raise CliError('--output needs exactly one input image. Use --output-dir for folders or multiple images.')
    targets = set()
    result = []
    for source, relative in found.items():
        target = (args.output or output_root / relative.parent / (relative.stem + '_nobg.png')).absolute()
        resolved = target.resolve()
        target_id = str(resolved).casefold()
        if target.suffix.lower() != '.png':
            raise CliError('Output files must use the .png extension.')
        if any(target_id == str(src).casefold() or (target.exists() and os.path.samefile(target, src)) for src in found):
            raise CliError(f'Output would overwrite an input: {target}')
        if target_id in targets:
            raise CliError(f'Output name collision: {target}. Process these files separately with distinct --output names.')
        if target.is_symlink() or (target.exists() and not target.is_file()):
            raise CliError(f'Output is not a regular file: {target}')
        if target.exists() and not args.overwrite:
            raise CliError(f'Output exists: {target}. Choose another path or use --overwrite.')
        size = source.stat().st_size
        if not 0 < size <= MAX_INPUT:
            raise CliError(f'Input must be between 1 byte and 25 MiB: {source}')
        targets.add(target_id)
        result.append((source, target))
    return result


def multipart(source, size):
    # Read with a bound even if the file changes after preflight.
    with source.open('rb') as f:
        content = f.read(MAX_INPUT + 1)
    if not 0 < len(content) <= MAX_INPUT:
        raise CliError('Input is empty or exceeds 25 MiB.')
    boundary = 'bgclear-' + uuid.uuid4().hex
    fields = {'size': size, 'format': 'png', 'bg_color': 'transparent', 'crop': 'false', 'store': 'false'}
    chunks = [f'--{boundary}\r\nContent-Disposition: form-data; name="{k}"\r\n\r\n{v}\r\n'.encode() for k, v in fields.items()]
    # Fixed filename avoids leaking local names or injecting multipart headers.
    chunks.append(f'--{boundary}\r\nContent-Disposition: form-data; name="image_file"; filename="image{source.suffix.lower()}"\r\nContent-Type: application/octet-stream\r\n\r\n'.encode())
    chunks.extend([content, f'\r\n--{boundary}--\r\n'.encode()])
    return b''.join(chunks), boundary


def retry_delay(header, attempt):
    if not header:
        return float(10 * (attempt + 1))
    try:
        delay = float(header)
    except ValueError:
        try:
            delay = email.utils.parsedate_to_datetime(header).timestamp() - time.time()
        except (ValueError, TypeError, OverflowError):
            return float(10 * (attempt + 1))
    if not math.isfinite(delay) or delay > 120:
        raise CliError('Server requested a long retry delay. Try again later; no further automatic retries will be made.', fatal=True)
    return max(1.0, delay)


def validate_png(data):
    if not data.startswith(b'\x89PNG\r\n\x1a\n'):
        raise CliError('Server did not return a PNG; no output was saved.')
    pos = 8
    saw_header = saw_data = False
    while pos + 12 <= len(data):
        length = struct.unpack('>I', data[pos:pos + 4])[0]
        kind = data[pos + 4:pos + 8]
        end = pos + 12 + length
        if end > len(data):
            break
        chunk = data[pos + 8:pos + 8 + length]
        crc = struct.unpack('>I', data[end - 4:end])[0]
        if zlib.crc32(kind + chunk) & 0xffffffff != crc:
            raise CliError('PNG integrity check failed; no output was saved.')
        if not saw_header:
            if kind != b'IHDR' or length != 13 or not all(struct.unpack('>II', chunk[:8])):
                break
            saw_header = True
        elif kind == b'IDAT':
            saw_data = True
        elif kind == b'IEND':
            if saw_data and length == 0 and end == len(data):
                return
            break
        pos = end
    raise CliError('Incomplete or invalid PNG; no output was saved.')


def remove(source, size, key):
    body, boundary = multipart(source, size)
    request = urllib.request.Request(API_URL, data=body, headers={
        'Authorization': 'Bearer ' + key,
        'User-Agent': 'BGClearCLI/' + VERSION,
        'Content-Type': 'multipart/form-data; boundary=' + boundary,
        'Accept': 'image/png',
        'Idempotency-Key': str(uuid.uuid4()),
    })
    opener = urllib.request.build_opener(NoRedirect())
    for attempt in range(4):
        try:
            with opener.open(request, timeout=90) as response:
                data = response.read(MAX_OUTPUT + 1)
            if len(data) > MAX_OUTPUT:
                raise CliError('Response exceeds the 48 MiB safety limit; no output was saved.')
            validate_png(data)
            return data
        except urllib.error.HTTPError as error:
            with error:
                raw = error.read(65536)
                try:
                    detail = json.loads(raw)
                    detail = detail.get('detail', detail)
                    code = detail.get('error', {}).get('code', '') if isinstance(detail, dict) else ''
                except (ValueError, AttributeError):
                    code = ''
                retryable = error.code == 429 or (error.code == 503 and code == 'gpu_busy')
                if retryable and attempt < 3:
                    delay = retry_delay(error.headers.get('Retry-After'), attempt)
                    print(f'Service busy (HTTP {error.code}); retrying in {delay:g}s ({attempt + 1}/3).', file=sys.stderr)
                    time.sleep(delay)
                    continue
                tips = {
                    400: 'Check the image. Full resolution supports inputs up to 4 MP.',
                    401: 'Check BGCLEAR_API_KEY or create a new key.',
                    402: 'No full-resolution credits available. Use --size preview for free previews.',
                    403: 'Access denied. Preview-only keys cannot process --size full; check your key scope.',
                    409: 'This request cannot be replayed. Check usage before starting a new full-resolution request.',
                    413: 'The upload is too large.',
                    429: 'Rate limit reached. Wait before running again.',
                    503: 'Processing is unavailable. Try again later.',
                }
                hint = tips.get(error.code, 'See https://www.bgclear.ai/api-docs/#errors for help.')
                raise CliError(f'HTTP {error.code}. {hint}', fatal=error.code in (401, 402, 403, 429) or retryable) from None
        except (urllib.error.URLError, TimeoutError, http.client.HTTPException, ConnectionError):
            raise CliError('Connection failed or response was interrupted. Not retried automatically: processing may have completed. Check API usage before retrying full resolution.', fatal=True) from None
    raise CliError('Retry limit reached.', fatal=True)


def save(data, target, overwrite, sources):
    target.parent.mkdir(parents=True, exist_ok=True)
    # Recheck after the request: paths may have changed while processing.
    if target.is_symlink() or any(target.resolve() == src or (target.exists() and os.path.samefile(target, src)) for src in sources):
        raise CliError('Output path changed to an input or symbolic link; refusing to replace it.')
    fd, temp = tempfile.mkstemp(prefix='.bgclear-', suffix='.tmp', dir=target.parent)
    try:
        with os.fdopen(fd, 'wb') as f:
            f.write(data)
            f.flush()
            os.fsync(f.fileno())
        if overwrite:
            os.replace(temp, target)
        else:
            # Atomic no-clobber publish. Supported on NTFS, APFS, ext4; fails safely
            # on filesystems without hard links rather than overwriting a new file.
            os.link(temp, target)
    finally:
        Path(temp).unlink(missing_ok=True)


def main(argv=None):
    args = parser().parse_args(argv)
    try:
        key = os.environ.get('BGCLEAR_API_KEY', '').strip()
        if not key or any(ord(c) < 33 or ord(c) > 126 for c in key):
            raise CliError('Set BGCLEAR_API_KEY to your API key. Create a preview-only key at https://www.bgclear.ai/api-keys/ to prevent paid processing.')
        tasks = plan(args)
        if args.size == 'full':
            print('Full resolution selected: 1 credit per successful image; maximum 4 MP input.', file=sys.stderr)
        completed = 0
        sources = {source for source, _ in tasks}
        for index, (source, target) in enumerate(tasks):
            if index:
                time.sleep(6.1)
            try:
                data = remove(source, args.size, key)
                save(data, target, args.overwrite, sources)
                completed += 1
                print(f'Saved: {target} ({len(data):,} bytes)')
            except (CliError, OSError) as error:
                # Do not print request objects, server bodies, or secrets.
                message = str(error).replace(key, '[redacted]') if isinstance(error, CliError) else 'Could not read or save a file. Check permissions, free disk space and output paths. Processing may already have completed.'
                print(f'Failed: {source.name}: {message}', file=sys.stderr)
                if isinstance(error, OSError) or error.fatal:
                    break
        print(f'{completed}/{len(tasks)} images saved.')
        return 0 if completed == len(tasks) else 1
    except (CliError, OSError) as error:
        message = str(error).replace(key, '[redacted]') if isinstance(error, CliError) and key else (str(error) if isinstance(error, CliError) else 'Could not inspect inputs or outputs. Check paths and permissions.')
        print(f'BGClear: {message}', file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        print('\nCancelled. A request already sent may have completed; check usage before retrying full resolution.', file=sys.stderr)
        return 130


if __name__ == '__main__':
    sys.exit(main())