#!/usr/bin/env python3
"""
Atlas API Client - Cross-platform script for document extraction
Works on Windows, Linux, and macOS
"""

import os
import sys
import json
import time
import subprocess
import shutil
from pathlib import Path
try:
    import requests
except ImportError:
    print("❌ Error: 'requests' library not found.")
    print("📦 Install it with: pip install requests")
    sys.exit(1)


def clear_screen():
    """Clear terminal screen (cross-platform)"""
    os.system('cls' if os.name == 'nt' else 'clear')


def print_header():
    """Print script header"""
    print("╔════════════════════════════════════════╗")
    print("║      Atlas API Client - v1.0           ║")
    print("║   Document Extraction & Processing     ║")
    print("╚════════════════════════════════════════╝")
    print()


def get_user_input(prompt, default=None):
    """Get user input with optional default value"""
    if default:
        user_input = input(f"{prompt} [{default}]: ").strip()
        return user_input if user_input else default
    return input(f"{prompt}: ").strip()


def list_operations(base_url, api_key):
    """Fetch and display available operations"""
    try:
        print("\n🔍 Fetching available operations...")
        response = requests.get(
            f"{base_url}/api/operations",
            headers={"x-api-key": api_key},
            timeout=10
        )

        if response.status_code == 200:
            data = response.json()
            operations = data.get("operations", [])

            if not operations:
                print("⚠️  No operations available")
                return []

            print("\n📋 Available Operations:")
            print("-" * 70)
            for idx, op in enumerate(operations, 1):
                cost = op.get('cost', 0)
                code = op.get('code', 'unknown')
                name = op.get('name', 'Unknown')
                print(f"{idx}. {name}")
                print(f"   Code: {code}")
                print(f"   Cost: {cost} credits")
                print()

            return operations
        else:
            print(f"❌ Error fetching operations: {response.status_code}")
            print(f"   Response: {response.text}")
            return []

    except requests.exceptions.RequestException as e:
        print(f"❌ Connection error: {e}")
        return []


def detect_file_format(file_path):
    """Detect actual file format by reading magic bytes. Returns: 'pdf', 'tiff', 'png', 'jpeg', or None."""
    try:
        with open(file_path, 'rb') as f:
            header = f.read(16)
        if len(header) < 4:
            return None
        # PDF: %PDF
        if header[:4] == b'%PDF':
            return 'pdf'
        # TIFF classic: II*\x00 (little-endian) or MM\x00* (big-endian)
        if header[:4] in (b'II\x2a\x00', b'MM\x00\x2a'):
            return 'tiff'
        # BigTIFF: II+\x00 (little-endian) or MM\x00+ (big-endian)
        if header[:4] in (b'II\x2b\x00', b'MM\x00\x2b'):
            return 'tiff'
        # PNG: \x89PNG\r\n\x1a\n
        if header[:8] == b'\x89PNG\r\n\x1a\n':
            return 'png'
        # JPEG: \xff\xd8\xff
        if header[:3] == b'\xff\xd8\xff':
            return 'jpeg'
        # BMP: BM
        if header[:2] == b'BM':
            return 'bmp'
        # Print raw bytes for debugging unrecognized files
        hex_str = ' '.join(f'{b:02x}' for b in header[:16])
        print(f"\n   DEBUG magic bytes of {file_path}: {hex_str}")
        return None
    except (OSError, IOError):
        return None


def find_files_in_directory(directory=".", extensions=None, depth=1):
    """Find files with given extensions in directory (depth 0 = current only, 1 = include subdirs)"""
    if extensions is None:
        extensions = [".pdf"]
    path = Path(directory)
    files = []
    for ext in extensions:
        # depth 0: current dir only
        files.extend(path.glob(f"*{ext}"))
        files.extend(path.glob(f"*{ext.upper()}"))
        # depth 1: one level of subdirectories
        if depth >= 1:
            files.extend(path.glob(f"*/*{ext}"))
            files.extend(path.glob(f"*/*{ext.upper()}"))
    # Exclude files inside atlas_converted directory
    files = [f for f in files if "atlas_converted" not in f.parts]
    return sorted(set(files))


def find_conversion_tools():
    """Find all available TIFF-to-PDF conversion tools, ordered by preference."""
    tools = []
    # On Windows, 'convert' is a system utility (FAT->NTFS), NOT ImageMagick
    if shutil.which("magick"):
        tools.append("magick")
    if shutil.which("gdal_translate"):
        tools.append("gdal_translate")
    if shutil.which("ffmpeg"):
        tools.append("ffmpeg")
    if shutil.which("tiff2pdf"):
        tools.append("tiff2pdf")
    # Only trust 'convert' on non-Windows systems
    if os.name != 'nt' and shutil.which("convert"):
        tools.append("convert")
    # Pillow
    try:
        from PIL import Image
        tools.append("pillow")
    except ImportError:
        pass
    # Raw TIFF reader fallback (pure Python, no external deps beyond Pillow for PDF output)
    try:
        from PIL import Image
        tools.append("raw")
    except ImportError:
        pass
    return tools


def convert_tiff_pillow(tiff_path, pdf_path):
    """Convert a TIFF to PDF using Pillow (pure Python fallback)."""
    from PIL import Image
    img = Image.open(tiff_path)
    images = []
    # Handle multi-page TIFFs
    try:
        while True:
            images.append(img.copy().convert('RGB'))
            img.seek(img.tell() + 1)
    except EOFError:
        pass
    if images:
        images[0].save(pdf_path, "PDF", save_all=True, append_images=images[1:])


def convert_tiff_raw(tiff_path, pdf_path):
    """Convert TIFF to PDF by parsing the TIFF structure and extracting embedded image data.
    Handles JPEG-in-TIFF (compression 6/7), which Pillow cannot open directly."""
    import struct
    import io
    from PIL import Image

    abs_tiff = str(Path(tiff_path).resolve())
    abs_pdf = str(Path(pdf_path).resolve())

    with open(abs_tiff, 'rb') as f:
        data = f.read()

    # Determine byte order
    if data[:2] == b'II':
        endian = '<'
    elif data[:2] == b'MM':
        endian = '>'
    else:
        raise ValueError("Not a TIFF file")

    ifd_offset = struct.unpack_from(endian + 'I', data, 4)[0]

    def read_ifd(offset):
        num_entries = struct.unpack_from(endian + 'H', data, offset)[0]
        tags = {}
        pos = offset + 2
        for _ in range(num_entries):
            tag_id, tag_type, count, value_offset = struct.unpack_from(endian + 'HHII', data, pos)
            if tag_type == 3 and count == 1:  # SHORT
                value = struct.unpack_from(endian + 'H', data, pos + 8)[0]
            elif tag_type == 4 and count == 1:  # LONG
                value = value_offset
            else:
                value = value_offset
            tags[tag_id] = (tag_type, count, value, value_offset)
            pos += 12
        next_ifd = struct.unpack_from(endian + 'I', data, pos)[0]
        return tags, next_ifd

    def get_tag_value(tags, tag_id, default=None):
        if tag_id not in tags:
            return default
        return tags[tag_id][2]

    def read_offsets(tag_entry):
        """Read an array of LONG values from a tag (strip/tile offsets or byte counts)."""
        _, count, value, value_offset = tag_entry
        if count == 1:
            return [value_offset]
        fmt = endian + str(count) + 'I'
        return list(struct.unpack_from(fmt, data, value_offset))

    # Collect all IFD offsets (multi-page)
    ifd_offsets = []
    offset = ifd_offset
    while offset != 0:
        ifd_offsets.append(offset)
        _, next_off = read_ifd(offset)
        offset = next_off

    images = []
    for ifd_off in ifd_offsets:
        tags, _ = read_ifd(ifd_off)
        compression = get_tag_value(tags, 259, 1)

        # --- JPEG-in-TIFF (compression 6 = old JPEG, 7 = new JPEG) ---
        if compression in (6, 7):
            # Try strips first, then tiles
            strip_tag = tags.get(273)   # StripOffsets
            count_tag = tags.get(279)   # StripByteCounts
            tile_tag = tags.get(324)    # TileOffsets
            tile_count_tag = tags.get(325)  # TileByteCounts

            if strip_tag and count_tag:
                offsets = read_offsets(strip_tag)
                counts = read_offsets(count_tag)
            elif tile_tag and tile_count_tag:
                offsets = read_offsets(tile_tag)
                counts = read_offsets(tile_count_tag)
            else:
                continue

            # For single-strip JPEG, the strip IS the full JPEG
            if len(offsets) == 1:
                jpeg_data = data[offsets[0]:offsets[0] + counts[0]]
                # Verify it starts with JPEG SOI marker
                if jpeg_data[:2] == b'\xff\xd8':
                    img = Image.open(io.BytesIO(jpeg_data))
                    images.append(img.convert('RGB'))
                    continue

            # Multi-strip JPEG: concatenate all strips
            jpeg_data = b''
            for so, sc in zip(offsets, counts):
                jpeg_data += data[so:so + sc]

            if jpeg_data[:2] == b'\xff\xd8':
                img = Image.open(io.BytesIO(jpeg_data))
                images.append(img.convert('RGB'))
                continue

            # If JPEGInterchangeFormat tag exists (tag 513), use that
            jif_offset = get_tag_value(tags, 513)
            jif_length = get_tag_value(tags, 514)
            if jif_offset and jif_length:
                jpeg_data = data[jif_offset:jif_offset + jif_length]
                if jpeg_data[:2] == b'\xff\xd8':
                    img = Image.open(io.BytesIO(jpeg_data))
                    images.append(img.convert('RGB'))
                    continue

            # Last resort: scan for JPEG SOI/EOI in the file around the IFD area
            # Skip this page
            continue

        # --- Uncompressed or simple compressions ---
        else:
            width = get_tag_value(tags, 256, 0)
            height = get_tag_value(tags, 257, 0)
            bits = get_tag_value(tags, 258, 8)
            photometric = get_tag_value(tags, 262, 0)
            samples = get_tag_value(tags, 277, 1)

            strip_tag = tags.get(273)
            count_tag = tags.get(279)
            if not strip_tag or not count_tag or width == 0 or height == 0:
                continue

            offsets = read_offsets(strip_tag)
            counts = read_offsets(count_tag)

            raw_data = b''
            for so, sc in zip(offsets, counts):
                raw_data += data[so:so + sc]

            if compression == 1:  # Uncompressed
                pixel_data = raw_data
            elif compression in (8, 32946):  # Deflate / Adobe Deflate
                import zlib
                pixel_data = zlib.decompress(raw_data)
            elif compression == 32773:  # PackBits
                pixel_data = _unpack_packbits(raw_data)
            elif compression == 5:  # LZW
                pixel_data = _decode_lzw(raw_data)
            else:
                raise ValueError(f"Unsupported compression: {compression}")

            if samples >= 3:
                mode = 'RGB'
                expected = width * height * 3
            elif bits == 1:
                mode = '1'
                expected = (width + 7) // 8 * height
            else:
                mode = 'L'
                expected = width * height

            if len(pixel_data) < expected:
                pixel_data += b'\x00' * (expected - len(pixel_data))

            if mode == '1':
                img = Image.frombytes('1', (width, height), pixel_data[:expected], 'raw', '1', 0, 1)
            else:
                img = Image.frombytes(mode, (width, height), pixel_data[:expected])

            if photometric == 0 and mode == 'L':
                from PIL import ImageOps
                img = ImageOps.invert(img)

            images.append(img.convert('RGB'))

    if not images:
        raise RuntimeError("Could not extract any pages from TIFF")

    images[0].save(abs_pdf, "PDF", save_all=True, append_images=images[1:])
    for img in images:
        img.close()


def _unpack_packbits(data):
    """Decode PackBits compressed data."""
    result = bytearray()
    i = 0
    while i < len(data):
        n = data[i]
        if n < 128:
            count = n + 1
            result.extend(data[i + 1:i + 1 + count])
            i += 1 + count
        elif n > 128:
            count = 257 - n
            result.extend(bytes([data[i + 1]]) * count)
            i += 2
        else:
            i += 1
    return bytes(result)


def _decode_lzw(data):
    """Decode LZW compressed TIFF data."""
    result = bytearray()
    CLEAR = 256
    EOI = 257
    table = {i: bytes([i]) for i in range(256)}
    table[CLEAR] = b''
    table[EOI] = b''
    next_code = 258
    code_size = 9
    bit_pos = 0
    prev = None

    def read_code():
        nonlocal bit_pos
        byte_pos = bit_pos >> 3
        bit_offset = bit_pos & 7
        # Read up to 4 bytes
        raw = 0
        for j in range(4):
            if byte_pos + j < len(data):
                raw |= data[byte_pos + j] << (24 - j * 8)
        code = (raw >> (32 - bit_offset - code_size)) & ((1 << code_size) - 1)
        bit_pos += code_size
        return code

    while bit_pos < len(data) * 8:
        code = read_code()
        if code == EOI:
            break
        if code == CLEAR:
            table = {i: bytes([i]) for i in range(256)}
            table[CLEAR] = b''
            table[EOI] = b''
            next_code = 258
            code_size = 9
            prev = None
            continue
        if code in table:
            entry = table[code]
        elif code == next_code and prev is not None:
            entry = prev + prev[:1]
        else:
            break
        result.extend(entry)
        if prev is not None and next_code < 4096:
            table[next_code] = prev + entry[:1]
            next_code += 1
            if next_code > (1 << code_size) - 1 and code_size < 12:
                code_size += 1
        prev = entry

    return bytes(result)


def try_convert_single(tiff_path, pdf_path, tools):
    """Try converting a single TIFF using available tools in order. Returns (success, error_msg)."""
    last_err = "no tools available"
    for tool in tools:
        try:
            if tool == "pillow":
                convert_tiff_pillow(tiff_path, pdf_path)
            elif tool == "raw":
                convert_tiff_raw(tiff_path, pdf_path)
            elif tool == "tiff2pdf":
                subprocess.run(
                    ["tiff2pdf", "-o", str(pdf_path), str(tiff_path)],
                    check=True, capture_output=True
                )
            elif tool == "magick":
                subprocess.run(
                    ["magick", "convert", str(tiff_path), str(pdf_path)],
                    check=True, capture_output=True
                )
            elif tool == "gdal_translate":
                subprocess.run(
                    ["gdal_translate", "-of", "PDF", str(tiff_path), str(pdf_path)],
                    check=True, capture_output=True
                )
            elif tool == "ffmpeg":
                subprocess.run(
                    ["ffmpeg", "-y", "-i", str(tiff_path), str(pdf_path)],
                    check=True, capture_output=True
                )
            else:  # convert (non-Windows only)
                subprocess.run(
                    ["convert", str(tiff_path), str(pdf_path)],
                    check=True, capture_output=True
                )
            # Verify output exists and is non-empty
            if pdf_path.exists() and pdf_path.stat().st_size > 0:
                return True, None
            else:
                last_err = f"{tool}: produced empty output"
                if pdf_path.exists():
                    pdf_path.unlink()
        except Exception as e:
            # Clean up partial file before trying next tool
            if pdf_path.exists():
                pdf_path.unlink()
            try:
                last_err = f"{tool}: {e}"
                if hasattr(e, 'stderr') and e.stderr:
                    last_err = f"{tool}: {e.stderr.decode(errors='replace').strip()}"
            except Exception:
                last_err = f"{tool}: {repr(e)}"
            continue
    return False, last_err


def convert_tiffs_to_pdf(tiff_files, output_dir="atlas_converted"):
    """Convert TIFF files to PDF in a separate output folder. Originals are not modified."""
    converted = []
    skipped = 0
    failed = 0
    not_tiff = 0

    tools = find_conversion_tools()
    if not tools:
        print("❌ No TIFF conversion tool found.")
        print("📦 Install one of:")
        if os.name == 'nt':
            print("   - Pillow (easiest): pip install Pillow")
            print("   - ImageMagick: https://imagemagick.org/script/download.php#windows")
        else:
            print("   - Pillow (easiest): pip install Pillow")
            print("   - ImageMagick: sudo apt install imagemagick (Linux) / brew install imagemagick (macOS)")
            print("   - libtiff:     sudo apt install libtiff-tools (Linux) / brew install libtiff (macOS)")
        return []

    print(f"   Available conversion tools: {', '.join(tools)}")

    # Create output directory
    out_path = Path(output_dir)
    out_path.mkdir(exist_ok=True)

    for tiff_path in tiff_files:
        # Check magic bytes (informational - still attempt conversion)
        actual_format = detect_file_format(tiff_path)
        if actual_format is not None and actual_format != 'tiff':
            # Definitely a different format (PDF, PNG, JPEG, BMP) - skip
            not_tiff += 1
            print(f"   ⚠️  Skipping {tiff_path} (actually {actual_format}, not TIFF)")
            continue
        if actual_format is None:
            print(f"   ℹ️  {tiff_path} has unknown magic bytes, will attempt conversion anyway")

        # Build output path preserving subdirectory structure
        relative = tiff_path.parent
        dest_dir = out_path / relative
        dest_dir.mkdir(parents=True, exist_ok=True)
        pdf_path = dest_dir / tiff_path.with_suffix(".pdf").name

        # Skip already-converted files
        if pdf_path.exists() and pdf_path.stat().st_size > 0:
            skipped += 1
            converted.append(pdf_path)
            continue

        total = len(tiff_files) - not_tiff
        current = len(converted) + failed + 1
        print(f"   🔄 [{current}/{total}] {tiff_path} ...", end=" ", flush=True)

        success, err_msg = try_convert_single(tiff_path, pdf_path, tools)
        if success:
            print("✅")
            converted.append(pdf_path)
        else:
            failed += 1
            print(f"❌ {err_msg}")

    if skipped:
        print(f"\n   ⏭️  {skipped} file(s) already converted (skipped)")
    if not_tiff:
        print(f"   ⚠️  {not_tiff} file(s) not actually TIFF format (skipped)")
    if failed:
        print(f"   ❌ {failed} file(s) failed to convert")
    print(f"   ✅ {len(converted)} file(s) ready")

    return converted


def send_document(base_url, api_key, file_path, operation_code, await_result=False):
    """Send a document for processing"""
    try:
        print(f"\n📤 Uploading: {file_path.name}")

        mime_type = 'application/pdf' if file_path.suffix.lower() == '.pdf' else f'image/{file_path.suffix.lower().lstrip(".")}'
        with open(file_path, 'rb') as f:
            files = {'file': (file_path.name, f, mime_type)}
            data = {
                'operation': operation_code,
                'awaitResult': 'true' if await_result else 'false'
            }

            response = requests.post(
                f"{base_url}/api/operations",
                headers={"x-api-key": api_key},
                files=files,
                data=data,
                timeout=300 if await_result else 30
            )

        if response.status_code in [200, 202]:
            result = response.json()
            task_id = result.get('taskId', 'unknown')
            status = result.get('status', 'unknown')

            print(f"✅ Success! Task ID: {task_id}")
            print(f"   Status: {status}")

            if status == 'completed' and 'result' in result:
                print(f"   Result available!")
                return {'success': True, 'taskId': task_id, 'result': result}
            else:
                print(f"   Check status at: {base_url}/api/operations/{task_id}")
                return {'success': True, 'taskId': task_id, 'result': None}

        elif response.status_code == 402:
            print(f"❌ Insufficient credits")
            return {'success': False, 'error': 'insufficient_credits'}

        else:
            print(f"❌ Error: {response.status_code}")
            print(f"   Response: {response.text}")
            return {'success': False, 'error': response.text}

    except requests.exceptions.RequestException as e:
        print(f"❌ Connection error: {e}")
        return {'success': False, 'error': str(e)}

    except Exception as e:
        print(f"❌ Unexpected error: {e}")
        return {'success': False, 'error': str(e)}


def check_task_status(base_url, api_key, task_id):
    """Check the status of a task"""
    try:
        response = requests.get(
            f"{base_url}/api/operations/{task_id}",
            headers={"x-api-key": api_key},
            timeout=10
        )

        if response.status_code == 200:
            return response.json()
        else:
            print(f"❌ Error checking status: {response.status_code}")
            return None

    except requests.exceptions.RequestException as e:
        print(f"❌ Connection error: {e}")
        return None


def save_results(results, output_dir="atlas_results"):
    """Save results to JSON files"""
    output_path = Path(output_dir)
    output_path.mkdir(exist_ok=True)

    timestamp = time.strftime("%Y%m%d_%H%M%S")
    filename = output_path / f"results_{timestamp}.json"

    with open(filename, 'w', encoding='utf-8') as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

    print(f"\n💾 Results saved to: {filename}")


def main():
    """Main script execution"""
    clear_screen()
    print_header()

    # Get API configuration
    print("🔧 Configuration")
    print("-" * 70)
    base_url = get_user_input("API Base URL", "https://backend-atlas.registrotecnologia.com.br")

    # Remove trailing /api/operations if user included it
    base_url = base_url.rstrip('/')
    if base_url.endswith('/api/operations'):
        base_url = base_url[:-len('/api/operations')]

    api_key = get_user_input("API Key")

    if not api_key:
        print("❌ API key is required!")
        sys.exit(1)

    # List available operations
    operations = list_operations(base_url, api_key)

    if not operations:
        print("\n❌ Cannot proceed without available operations")
        sys.exit(1)

    # Select operation
    print("-" * 70)
    while True:
        try:
            choice = int(get_user_input("\nSelect operation number"))
            if 1 <= choice <= len(operations):
                selected_operation = operations[choice - 1]
                break
            else:
                print(f"⚠️  Please enter a number between 1 and {len(operations)}")
        except ValueError:
            print("⚠️  Please enter a valid number")

    operation_code = selected_operation['code']
    operation_name = selected_operation['name']
    print(f"\n✓ Selected: {operation_name} ({operation_code})")

    # Find PDFs and PNGs in current directory (depth 1)
    pdfs = find_files_in_directory(".", extensions=[".pdf", ".png"], depth=1)

    if not pdfs:
        # No PDFs or PNGs - check for TIFFs and convert
        tiffs = find_files_in_directory(".", extensions=[".tif", ".tiff"], depth=1)
        if tiffs:
            print(f"\n⚠️  No PDF/PNG files found, but found {len(tiffs)} TIFF(s)")
            print("🔄 Converting TIFFs to PDF...")
            converted = convert_tiffs_to_pdf(tiffs)
            if converted:
                print(f"\n✅ Converted {len(converted)} file(s)")
                pdfs = converted
            else:
                print("\n❌ No files could be converted")

    if pdfs:
        print(f"\n📁 Found {len(pdfs)} file(s):")
        for pdf in pdfs[:10]:  # Show first 10
            relative = pdf if pdf.parent == Path(".") else pdf
            print(f"   - {relative}")
        if len(pdfs) > 10:
            print(f"   ... and {len(pdfs) - 10} more")

        print()
        filter_recent = get_user_input("Filter only recently created files? (y/n)", "n").lower() == 'y'
        if filter_recent:
            hours = float(get_user_input("Created in the last N hours?", "12"))
            import datetime
            cutoff = datetime.datetime.now().timestamp() - (hours * 3600)
            before = len(pdfs)
            pdfs = [p for p in pdfs if p.stat().st_ctime >= cutoff]
            print(f"   Filtered: {before} -> {len(pdfs)} file(s) created in the last {hours}h")

        process_all = get_user_input("Process all files? (y/n)", "n").lower() == 'y' if pdfs else False
    else:
        print("\n⚠️  No processable files found in current directory")
        process_all = False

    # Ask if user wants to await results
    await_result = get_user_input("\nWait for results? (y/n)", "n").lower() == 'y'

    # Process files
    results = []

    if process_all and pdfs:
        print(f"\n🚀 Processing {len(pdfs)} file(s)...")
        print("=" * 70)

        for idx, doc in enumerate(pdfs, 1):
            print(f"\n[{idx}/{len(pdfs)}]")
            result = send_document(base_url, api_key, doc, operation_code, await_result)
            results.append({
                'file': str(doc),
                'taskId': result.get('taskId'),
                'success': result.get('success'),
                'result': result.get('result')
            })

            # Small delay between requests
            if idx < len(pdfs):
                time.sleep(0.5)

        print("\n" + "=" * 70)
        print(f"✅ Completed: {sum(1 for r in results if r['success'])}/{len(results)}")

        # Save results
        if results:
            save_results(results)

    else:
        # Single file upload
        file_path = get_user_input("\nEnter PDF file path")

        if not file_path:
            print("❌ File path is required!")
            sys.exit(1)

        pdf_path = Path(file_path)

        if not pdf_path.exists():
            print(f"❌ File not found: {file_path}")
            sys.exit(1)

        result = send_document(base_url, api_key, pdf_path, operation_code, await_result)

        if result['success'] and result['result']:
            print("\n📊 Result:")
            print(json.dumps(result['result'], indent=2))

            save_option = get_user_input("\nSave result to file? (y/n)", "y").lower()
            if save_option == 'y':
                save_results([{
                    'file': pdf_path.name,
                    'taskId': result.get('taskId'),
                    'success': True,
                    'result': result.get('result')
                }])

    print("\n✨ Done!")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n⚠️  Interrupted by user")
        sys.exit(0)
    except Exception as e:
        print(f"\n❌ Unexpected error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)