#!/usr/bin/env python3
"""
Sketchfab 3D Model Extractor

Extracts model files and textures from Sketchfab by parsing page metadata.
Decrypts .binz encrypted model files using embedded XOR key.

Usage:
  python sketchfab_extract.py <model_uid> <output_dir> [--timeout 120]
"""

import sys
import os
import json
import re
import html
import asyncio
import argparse
import gzip
import base64
import struct
from pathlib import Path
from urllib.parse import urlparse

try:
    from playwright.async_api import async_playwright, TimeoutError as PlaywrightTimeout
except ImportError:
    print(json.dumps({
        'success': False,
        'error': 'Playwright not installed. Run: pip install playwright && playwright install chromium'
    }))
    sys.exit(1)

try:
    import aiohttp
except ImportError:
    aiohttp = None


def output_json(data: dict):
    print(json.dumps(data))
    sys.exit(0 if data.get('success') else 1)


def sanitize_filename(name: str, max_length: int = 120) -> str:
    name = re.sub(r'[<>:"/\\\\|?*]', '_', name)
    name = re.sub(r'\\s+', '_', name).strip()
    name = re.sub(r'_+', '_', name)
    if len(name) > max_length:
        name = name[:max_length]
    return name or 'file'


def decrypt_binz(encrypted_data: bytes, key: bytes) -> bytes:
    """
    Decrypt Sketchfab .binz file using XOR with the provided key.
    The key is cycled through the entire file.
    """
    if not key:
        return encrypted_data

    key_len = len(key)
    decrypted = bytearray(len(encrypted_data))

    for i, byte in enumerate(encrypted_data):
        decrypted[i] = byte ^ key[i % key_len]

    return bytes(decrypted)


def decompress_osgjs(data: bytes) -> bytes:
    """
    Decompress gzip-compressed OSGJS data.
    """
    try:
        return gzip.decompress(data)
    except gzip.BadGzipFile:
        # Try without gzip header (raw deflate)
        import zlib
        try:
            return zlib.decompress(data, -zlib.MAX_WBITS)
        except:
            return data
    except Exception:
        return data


def extract_encryption_key(content: str) -> bytes:
    """
    Extract the encryption key from Sketchfab page content.
    The key is in JSON: "p": [{"v": 1, "b": "<base64 with literal \\n>", "d": true}]
    """
    # Pattern that matches base64 with literal \n escape sequences
    # In the HTML, it appears as backslash followed by 'n'
    match = re.search(r'"p":\s*\[\{"v":\s*\d+,\s*"b":\s*"([A-Za-z0-9+/=\\n]+)"', content)
    if match:
        try:
            # Remove literal backslash-n sequences from base64 string
            b64_clean = match.group(1).replace('\\n', '')
            return base64.b64decode(b64_clean)
        except Exception as e:
            print(f"Key decode failed: {e}", file=sys.stderr)

    return b''


def clean_embedded_url(url: str) -> str:
    if not url:
        return url
    url = url.replace('\\u002F', '/')
    url = url.replace('\\/', '/')
    return url


class SketchfabExtractor:
    def __init__(self, output_dir: str, timeout: int = 120, headed: bool = False, debug: bool = False, capture_network: bool = False):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.timeout = timeout
        self.headed = headed
        self.debug = debug
        self.capture_network = capture_network
        self.intercepted_model_data = None
        self.intercepted_bin_files = {}
        self.captured_network_files = []

    async def download_file(self, url: str, path: Path) -> bool:
        """Download a file from URL"""
        try:
            if aiohttp:
                async with aiohttp.ClientSession() as session:
                    async with session.get(url) as resp:
                        if resp.status == 200:
                            with open(path, 'wb') as f:
                                f.write(await resp.read())
                            return True
            else:
                import urllib.request
                urllib.request.urlretrieve(url, path)
                return True
        except Exception as e:
            print(f"Download failed: {e}", file=sys.stderr)
            return False

    async def intercept_model_response(self, route, request):
        """Intercept model file responses to capture decrypted data"""
        url = request.url
        response = await route.fetch()

        # Check if this is a model geometry file
        if '/files/' in url and ('.bin' in url or '.binz' in url):
            try:
                body = await response.body()
                file_id = url.split('/files/')[-1].split('/')[0]
                self.intercepted_bin_files[file_id] = body
                print(f"Intercepted bin file: {file_id} ({len(body)} bytes)", file=sys.stderr)
            except Exception as e:
                print(f"Interception error: {e}", file=sys.stderr)

        await route.fulfill(response=response)

    async def run(self, model_uid: str) -> dict:
        url = f'https://sketchfab.com/models/{model_uid}/embed'

        async with async_playwright() as p:
            browser = await p.chromium.launch(
                headless=not self.headed,
                args=[
                    '--no-sandbox',
                    '--disable-setuid-sandbox',
                    '--enable-webgl',
                    '--use-gl=swiftshader',
                ]
            )

            page = await browser.new_page()

            # Capture model/texture assets from network responses (optional)
            if self.capture_network:
                network_dir = self.output_dir / 'network'
                network_dir.mkdir(exist_ok=True)

                async def handle_response(response):
                    try:
                        url = response.url
                        if not (
                            re.search(r'\\.(gltf|glb|bin|binz|osgjs|ktx2|png|jpg|jpeg|webp)(\\?|$)', url, re.IGNORECASE)
                            or 'file.osgjs' in url
                            or 'model_file' in url
                        ):
                            return
                        body = await response.body()
                        if not body:
                            return
                        filename = sanitize_filename(url.split('/')[-1].split('?')[0])
                        out_path = network_dir / f'{len(self.captured_network_files):03d}_{filename}'
                        with open(out_path, 'wb') as f:
                            f.write(body)
                        self.captured_network_files.append(str(out_path))
                    except Exception:
                        pass

                page.on('response', handle_response)

            # Set up route interception for model files
            await page.route('**/*.bin*', self.intercept_model_response)
            await page.route('**/*.binz*', self.intercept_model_response)

            try:
                # Inject script to capture decrypted geometry before page loads
                await page.add_init_script("""
                    window.__grabbit = {
                        models: [],
                        textures: [],
                        textureCache: {},
                        textureDump: []
                    };

                    window.attachbody = function(obj) {
                        try {
                            if (!obj || obj._faked === true) return;
                            if (obj._name === "composer layer" || obj._name === "Ground - Geometry") return;
                            obj._faked = true;
                            window.__grabbit.models.push(obj);
                        } catch (e) {}
                    };

                    window.drawhookcanvas = function(e, imagemodel) {
                        try {
                            if (!imagemodel) return e;
                            if ((e.width === 128 && e.height === 128) || (e.width === 32 && e.height === 32) || (e.width === 64 && e.height === 64)) {
                                return e;
                            }

                            var alpha = e.options && e.options.format;
                            var filename_image = imagemodel.attributes && imagemodel.attributes.name;
                            var url_image = e.url;
                            var max_size = 0;
                            var obr = e;
                            if (imagemodel.attributes && imagemodel.attributes.images) {
                                imagemodel.attributes.images.forEach(function(img) {
                                    var alpha_ok = alpha === "A" ? (img.options && img.options.format === alpha) : true;
                                    var d = img.width;
                                    while (d % 2 === 0) d = d / 2;
                                    if (img.size > max_size && alpha_ok && d === 1) {
                                        max_size = img.size;
                                        url_image = img.url;
                                        obr = img;
                                    }
                                });
                            }
                            if (url_image && !window.__grabbit.textureCache[url_image]) {
                                window.__grabbit.textureCache[url_image] = filename_image || url_image.split('/').pop();
                            }
                            return obr;
                        } catch (e) {}
                        return e;
                    };

                    window.drawhookimg = function(gl, t) {
                        try {
                            var img = t[5];
                            if (!img) return;
                            var url = img.currentSrc || img.src;
                            var width = img.width || img.videoWidth;
                            var height = img.height || img.videoHeight;
                            if (!url || !width || !height) return;
                            if (!window.__grabbit.textureCache[url]) return;

                            var data = new Uint8Array(width * height * 4);
                            gl.readPixels(0, 0, width, height, gl.RGBA, gl.UNSIGNED_BYTE, data);

                            var halfHeight = height / 2 | 0;
                            var bytesPerRow = width * 4;
                            var temp = new Uint8Array(width * 4);
                            for (var y = 0; y < halfHeight; ++y) {
                                var topOffset = y * bytesPerRow;
                                var bottomOffset = (height - y - 1) * bytesPerRow;
                                temp.set(data.subarray(topOffset, topOffset + bytesPerRow));
                                data.copyWithin(topOffset, bottomOffset, bottomOffset + bytesPerRow);
                                data.set(temp, bottomOffset);
                            }

                            var canvas = document.createElement('canvas');
                            canvas.width = width;
                            canvas.height = height;
                            var context = canvas.getContext('2d');
                            var imageData = context.createImageData(width, height);
                            imageData.data.set(data);
                            context.putImageData(imageData, 0, 0);

                            var name = window.__grabbit.textureCache[url] || ('texture_' + window.__grabbit.textures.length);
                            var dataUrl = canvas.toDataURL('image/png');
                            window.__grabbit.textures.push({ name: name, dataUrl: dataUrl, width: width, height: height });
                        } catch (e) {}
                    };

                    // Texture Dumper (v5+) inspired WebGL hook
                    (function () {
                        const MIN_TEXTURE_WIDTH = 256;
                        const MIN_TEXTURE_HEIGHT = 256;
                        let webGLTextureIdx = 0;
                        let webGLTextureMap = {};
                        let originTextureArr = [];
                        let embedPrefetch = null;

                        function tryLoadPrefetched() {
                            try {
                                const node = document.querySelector("#js-dom-data-prefetched-data");
                                if (!node) return;
                                const raw = node.innerHTML.replace("\\x3C!--", "").replace("-->", "").replaceAll("&#34;", "\\"");
                                const prefetchedData = JSON.parse(raw);
                                embedPrefetch = prefetchedData;
                                const keys = Object.keys(prefetchedData || {});
                                const texKey = keys.find(k => k.includes("/textures?optimized=1"));
                                if (texKey && prefetchedData[texKey] && prefetchedData[texKey].results) {
                                    originTextureArr = prefetchedData[texKey].results;
                                }
                            } catch (e) {}
                        }

                        function getNameById(uid) {
                            const tex = originTextureArr.find(t => t.uid == uid);
                            return tex ? tex.name : null;
                        }

                        function parseTexImage2dArgs(args) {
                            let argments = {};
                            if (args.length == 6) {
                                let [target, level, internalformat, format, type, source] = args;
                                let { width, height, src } = source || {};
                                argments = { target, level, width, height, internalformat, format, type, src };
                            } else if (args.length == 9) {
                                let [target, level, internalformat, width, height, border, format, type, pixels] = args;
                                argments = { target, level, width, height, format, internalformat, type };
                            } else {
                                let [target] = args;
                                argments = { target };
                            }
                            return argments;
                        }

                        function hookWebGLRenderingContext(context) {
                            if (!context || !context.prototype) return;

                            const glCreateTexture = context.prototype.createTexture;
                            context.prototype.createTexture = function (...args) {
                                let texture = glCreateTexture.apply(this, args);
                                texture.__grabbit_name = webGLTextureIdx;
                                texture.__grabbit_gl = this;
                                webGLTextureMap[`${webGLTextureIdx}`] = texture;
                                webGLTextureIdx++;
                                return texture;
                            };

                            const glBindTexture = context.prototype.bindTexture;
                            context.prototype.bindTexture = function (...args) {
                                let target = args[0], texture = args[1];
                                if (texture) texture.__grabbit_target = target;
                                return glBindTexture.apply(this, args);
                            };

                            const glDeleteTexture = context.prototype.deleteTexture;
                            context.prototype.deleteTexture = function (...args) { return; };

                            const glTexImage2D = context.prototype.texImage2D;
                            let lstTexture;
                            context.prototype.texImage2D = function (...args) {
                                let texture = this.getParameter(this.TEXTURE_BINDING_2D) || this.getParameter(this.TEXTURE_BINDING_CUBE_MAP);
                                let argments = parseTexImage2dArgs(args);
                                if (texture && texture.__grabbit_target == argments.target) {
                                    texture.__grabbit_args = argments;
                                    let { width, height, src } = argments;
                                    if (
                                        width > MIN_TEXTURE_WIDTH &&
                                        height > MIN_TEXTURE_HEIGHT &&
                                        (width & width - 1) === 0 &&
                                        (height & height - 1) === 0 &&
                                        texture.__grabbit_target === this.TEXTURE_2D
                                    ) {
                                        if (src) {
                                            texture.__grabbit_src = src;
                                            if (lstTexture) lstTexture.__grabbit_lst = texture.__grabbit_name;
                                        }
                                        lstTexture = texture;
                                    }
                                }
                                return glTexImage2D.apply(this, args);
                            };
                        }

                        function readWebTextureData(texture) {
                            let { __grabbit_target, __grabbit_gl, __grabbit_args, __grabbit_src, __grabbit_lst } = texture;
                            if (!__grabbit_args || !__grabbit_gl) return;
                            let { level, width, height, format, type } = __grabbit_args;
                            if (
                                width > MIN_TEXTURE_WIDTH &&
                                height > MIN_TEXTURE_HEIGHT &&
                                (width & width - 1) === 0 &&
                                (height & height - 1) === 0 &&
                                __grabbit_target === __grabbit_gl.TEXTURE_2D &&
                                !__grabbit_src
                            ) {
                                let fb = __grabbit_gl.createFramebuffer();
                                __grabbit_gl.bindFramebuffer(__grabbit_gl.FRAMEBUFFER, fb);
                                __grabbit_gl.framebufferTexture2D(
                                    __grabbit_gl.FRAMEBUFFER,
                                    __grabbit_gl.COLOR_ATTACHMENT0,
                                    __grabbit_gl.TEXTURE_2D,
                                    texture,
                                    level
                                );
                                if (__grabbit_gl.checkFramebufferStatus(__grabbit_gl.FRAMEBUFFER) == __grabbit_gl.FRAMEBUFFER_COMPLETE) {
                                    let originName = null;
                                    let originTex = webGLTextureMap[__grabbit_lst];
                                    if (originTex && originTex.__grabbit_src) {
                                        let path = originTex.__grabbit_src.split("/");
                                        let uid = path[path.length - 2];
                                        originName = getNameById(uid);
                                    }
                                    if (!originName) return;
                                    let pixels = new Uint8Array(width * height * 4);
                                    __grabbit_gl.readPixels(0, 0, width, height, format, type, pixels);
                                    let canvas = document.createElement("canvas");
                                    canvas.width = width;
                                    canvas.height = height;
                                    let context2d = canvas.getContext("2d");
                                    let imageData = context2d.createImageData(width, height);
                                    imageData.data.set(pixels);
                                    context2d.putImageData(imageData, 0, 0);
                                    window.__grabbit.textureDump.push({
                                        name: originName,
                                        dataUrl: canvas.toDataURL("image/png"),
                                        width,
                                        height
                                    });
                                }
                            }
                        }

                        window.__grabbitDumpTextures = function () {
                            tryLoadPrefetched();
                            for (let key in webGLTextureMap) {
                                readWebTextureData(webGLTextureMap[key]);
                            }
                        };

                        document.addEventListener("DOMContentLoaded", function () {
                            tryLoadPrefetched();
                            let context = (embedPrefetch && embedPrefetch.embedOptions && embedPrefetch.embedOptions.webgl2) ? WebGL2RenderingContext : WebGLRenderingContext;
                            hookWebGLRenderingContext(context);
                        }, false);
                    })();

                    window.__captured_geometry = [];
                    window.__captured_osgjs = null;

                    // Hook XMLHttpRequest to capture responses
                    const originalXHROpen = XMLHttpRequest.prototype.open;
                    const originalXHRSend = XMLHttpRequest.prototype.send;

                    XMLHttpRequest.prototype.open = function(method, url, ...args) {
                        this._url = url;
                        return originalXHROpen.call(this, method, url, ...args);
                    };

                    XMLHttpRequest.prototype.send = function(...args) {
                        this.addEventListener('load', function() {
                            if (this._url && (this._url.includes('.bin') || this._url.includes('file.osgjs'))) {
                                try {
                                    const data = this.response;
                                    if (data && data.byteLength) {
                                        window.__captured_geometry.push({
                                            url: this._url,
                                            size: data.byteLength
                                        });
                                    }
                                } catch(e) {}
                            }
                        });
                        return originalXHRSend.call(this, ...args);
                    };

                    // Hook fetch as well
                    const originalFetch = window.fetch;
                    window.fetch = async function(url, ...args) {
                        const response = await originalFetch(url, ...args);
                        if (url && (url.includes('.bin') || url.includes('file.osgjs'))) {
                            try {
                                const clone = response.clone();
                                const data = await clone.arrayBuffer();
                                window.__captured_geometry.push({
                                    url: url,
                                    size: data.byteLength
                                });
                            } catch(e) {}
                        }
                        return response;
                    };
                """)

                # Patch Sketchfab viewer scripts (Fabulous-Ripper technique)
                async def patch_viewer_script(route, request):
                    url = request.url
                    if ("web/dist/" not in url and "standaloneViewer" not in url):
                        return await route.fallback()

                    response = await route.fetch()
                    text = await response.text()
                    original = text

                    def insert_after(regex, insert):
                        m = re.search(regex, text)
                        if not m:
                            return text, False
                        idx = m.end()
                        return text[:idx] + insert + text[idx:], True

                    text, ok0 = insert_after(r"A\\.renderInto\\(n,E,R\\)", ",i")
                    text, ok1 = insert_after(r"g\\.renderInto=function\\(e,i,r", ",image_data")
                    text, ok2 = insert_after(r"t\\.drawArrays\\(t\\.TRIANGLES,0,6\\)", ",window.drawhookimg(t,image_data)")
                    text, ok3 = insert_after(r"getResourceImage:function\\(e,t\\)\\{", "e = window.drawhookcanvas(e,this._imageModel);")
                    text, ok4 = insert_after(r"(this\\._stateCache\\.drawGeometry\\(this\\._graphicContext,t\\))", ";window.attachbody(t);")

                    if text != original and self.debug:
                        print(f"Patched viewer script: {url}", file=sys.stderr)
                        print(f"patches: renderInto1={ok0} renderInto2={ok1} drawArrays={ok2} getResourceImage={ok3} drawGeometry={ok4}", file=sys.stderr)

                    await route.fulfill(
                        status=response.status,
                        headers=response.headers,
                        body=text
                    )

                await page.route("**/*.js", patch_viewer_script)

                await page.goto(url, wait_until='domcontentloaded', timeout=60000)
                await page.wait_for_timeout(5000)  # Wait for model to load

                # Get page content and unescape HTML entities
                content = await page.content()
                content = html.unescape(content)

                # Check if geometry was captured
                captured_info = await page.evaluate('() => window.__captured_geometry || []')

                # Extract geometry via patched viewer globals
                extracted_models = await page.evaluate("""
                    () => {
                        const out = [];
                        if (!window.__grabbit || !window.__grabbit.models) return out;
                        let idx = 0;
                        for (const obj of window.__grabbit.models) {
                            try {
                                const list = [];
                                if (obj._primitives) {
                                    obj._primitives.forEach(function(p) {
                                        if (p && p.indices) {
                                            list.push({ mode: p.mode, indices: Array.from(p.indices._elements || []) });
                                        }
                                    });
                                }
                                const attr = obj._attributes || {};
                                const vertex = attr.Vertex ? Array.from(attr.Vertex._elements || []) : [];
                                const normal = attr.Normal ? Array.from(attr.Normal._elements || []) : [];
                                const uv = attr.TexCoord0 ? Array.from(attr.TexCoord0._elements || []) :
                                            attr.TexCoord1 ? Array.from(attr.TexCoord1._elements || []) :
                                            attr.TexCoord2 ? Array.from(attr.TexCoord2._elements || []) :
                                            attr.TexCoord3 ? Array.from(attr.TexCoord3._elements || []) : [];
                                out.push({
                                    name: obj._name || (obj._parents && obj._parents[0] && obj._parents[0]._name) || ('model_' + idx),
                                    vertex, normal, uv, primitives: list
                                });
                                idx++;
                            } catch (e) {}
                        }
                        return out;
                    }
                """)

                extracted_textures = await page.evaluate("""
                    () => (window.__grabbit && window.__grabbit.textures) ? window.__grabbit.textures : []
                """)

                # Trigger Texture Dumper style readback (if available)
                try:
                    await page.evaluate("() => window.__grabbitDumpTextures && window.__grabbitDumpTextures()")
                except Exception:
                    pass

                extracted_dump_textures = await page.evaluate("""
                    () => (window.__grabbit && window.__grabbit.textureDump) ? window.__grabbit.textureDump : []
                """)

                # Take screenshot
                screenshot_path = self.output_dir / 'screenshot.png'
                await page.screenshot(path=str(screenshot_path))

                await browser.close()

                # Save page content for debugging
                page_path = self.output_dir / 'page.html'
                with open(page_path, 'w', encoding='utf-8') as f:
                    f.write(content)

                files = [str(screenshot_path)]

                # Save any intercepted binary files
                for file_id, data in self.intercepted_bin_files.items():
                    bin_path = self.output_dir / f'geometry_{file_id[:16]}.bin'
                    with open(bin_path, 'wb') as f:
                        f.write(data)
                    files.append(str(bin_path))

                # Extract model name
                name_match = re.search(r'"name":\s*"([^"]+)"', content)
                model_name = name_match.group(1) if name_match else model_uid

                # Extract author
                author_match = re.search(r'"username":\s*"([^"]+)"', content)
                author = author_match.group(1) if author_match else ''

                # Find model file URL (.binz or .osgjs.gz)
                binz_match = re.search(r'"osgjsUrl":\s*"([^"]+\.binz)"', content)
                osgjs_match = re.search(r'"osgjsUrl":\s*"([^"]+\.osgjs(?:\.gz)?)"', content)
                bin_match = re.search(r'"url":\s*"([^"]+model_file[^"]*\.bin(?:\.gz)?)"', content)

                model_url = None
                if binz_match:
                    model_url = binz_match.group(1)
                elif osgjs_match:
                    model_url = osgjs_match.group(1)

                if model_url:
                    model_url = clean_embedded_url(model_url)

                # Try alternate file paths (old format)
                if model_url and '.binz' in model_url:
                    # Try to get unencrypted version
                    alt_url = model_url.replace('/file.binz', '/file.osgjs.gz')
                    # We'll try both

                # Find texture URLs with dimensions
                tex_pattern = r'"uid":\s*"([^"]+)"[^}]*"width":\s*(\d+)[^}]*"height":\s*(\d+)[^}]*"url":\s*"(https://media\.sketchfab\.com/models/[^"]+/textures/[^"]+)"'
                tex_matches = re.findall(tex_pattern, content)
                extra_urls = re.findall(r'(https://media\\.sketchfab\\.com/[^"\\s]+)', content)

                # Download textures (get largest versions)
                tex_dir = self.output_dir / 'textures'
                tex_dir.mkdir(exist_ok=True)

                # Group by texture UID prefix and get largest
                tex_by_group = {}
                for uid, w, h, url in tex_matches:
                    url = clean_embedded_url(url)
                    # Group by the texture folder (second to last path segment)
                    path_parts = urlparse(url).path.split('/')
                    group = path_parts[-2] if len(path_parts) > 2 else uid

                    size = int(w) * int(h)
                    if group not in tex_by_group or size > tex_by_group[group][0]:
                        tex_by_group[group] = (size, w, h, url)

                # Add any extra media.sketchfab.com URLs that look like textures
                for url in extra_urls:
                    url = clean_embedded_url(url)
                    if not re.search(r'\\.(png|jpg|jpeg|webp|ktx2)(\\?|$)', url, re.IGNORECASE):
                        continue
                    path_parts = urlparse(url).path.split('/')
                    group = path_parts[-2] if len(path_parts) > 2 else 'texture'
                    size = 0
                    if group not in tex_by_group:
                        tex_by_group[group] = (size, '0', '0', url)

                # Download largest version of each texture
                downloaded_textures = 0
                for group, (size, w, h, url) in tex_by_group.items():
                    ext = url.split('.')[-1]
                    tex_path = tex_dir / f'{group[:20]}_{w}x{h}.{ext}'

                    if await self.download_file(url, tex_path):
                        files.append(str(tex_path))
                        downloaded_textures += 1

                # Save extracted textures from WebGL hook
                webgl_textures = 0
                if extracted_textures:
                    tex_dir = self.output_dir / 'textures_webgl'
                    tex_dir.mkdir(exist_ok=True)
                    for idx, tex in enumerate(extracted_textures):
                        try:
                            data_url = tex.get('dataUrl')
                            if not data_url:
                                continue
                            name = sanitize_filename(tex.get('name') or f'texture_{idx}')
                            b64 = data_url.split(',')[1]
                            tex_path = tex_dir / f'{name}.png'
                            with open(tex_path, 'wb') as f:
                                f.write(base64.b64decode(b64))
                            files.append(str(tex_path))
                            webgl_textures += 1
                        except Exception:
                            pass

                # Save extracted textures from Texture Dumper hook
                dump_textures = 0
                if extracted_dump_textures:
                    tex_dir = self.output_dir / 'textures_dump'
                    tex_dir.mkdir(exist_ok=True)
                    for idx, tex in enumerate(extracted_dump_textures):
                        try:
                            data_url = tex.get('dataUrl')
                            if not data_url:
                                continue
                            name = sanitize_filename(tex.get('name') or f'dump_texture_{idx}')
                            b64 = data_url.split(',')[1]
                            tex_path = tex_dir / f'{name}.png'
                            with open(tex_path, 'wb') as f:
                                f.write(base64.b64decode(b64))
                            files.append(str(tex_path))
                            dump_textures += 1
                        except Exception:
                            pass

                # Save extracted geometry from viewer (OBJ)
                obj_files = 0
                if extracted_models:
                    for idx, model in enumerate(extracted_models):
                        try:
                            name = sanitize_filename(model.get('name') or f'model_{idx}')
                            obj_path = self.output_dir / f'{name}.obj'
                            with open(obj_path, 'w') as obj_file:
                                obj_file.write(f'o {name}\n')

                                vertices = model.get('vertex', [])
                                normals = model.get('normal', [])
                                uvs = model.get('uv', [])
                                primitives = model.get('primitives', [])

                                for i in range(0, len(vertices), 3):
                                    obj_file.write(f'v {vertices[i]} {vertices[i+1]} {vertices[i+2]}\n')
                                for i in range(0, len(normals), 3):
                                    obj_file.write(f'vn {normals[i]} {normals[i+1]} {normals[i+2]}\n')
                                for i in range(0, len(uvs), 2):
                                    obj_file.write(f'vt {uvs[i]} {uvs[i+1]}\n')

                                vn = len(normals) > 0
                                vt = len(uvs) > 0

                                for prim in primitives:
                                    mode = prim.get('mode')
                                    indices = prim.get('indices', [])
                                    if mode in (4, 5):
                                        strip = (mode == 5)
                                        j = 0
                                        while j + 2 < len(indices):
                                            order = [0, 1, 2]
                                            if strip and (j % 2 == 1):
                                                order = [0, 2, 1]
                                            face = []
                                            for k in range(3):
                                                face_num = indices[j + order[k]] + 1
                                                if vn or vt:
                                                    if vt and vn:
                                                        face.append(f"{face_num}/{face_num}/{face_num}")
                                                    elif vt:
                                                        face.append(f"{face_num}/{face_num}")
                                                    else:
                                                        face.append(f"{face_num}//{face_num}")
                                                else:
                                                    face.append(f"{face_num}")
                                            obj_file.write("f " + " ".join(face) + "\n")
                                            j = j + (1 if strip else 3)
                            files.append(str(obj_path))
                            obj_files += 1
                        except Exception:
                            pass

                # Extract encryption key for .binz decryption
                encryption_key = extract_encryption_key(content)

                # Download model file
                model_downloaded = False
                model_decrypted = False
                if model_url:
                    model_ext = 'binz' if '.binz' in model_url else 'osgjs.gz'
                    model_path = self.output_dir / f'model.{model_ext}'

                    if await self.download_file(model_url, model_path):
                        files.append(str(model_path))
                        model_downloaded = True

                        # Decrypt .binz files
                        if model_ext == 'binz' and encryption_key:
                            try:
                                with open(model_path, 'rb') as f:
                                    encrypted_data = f.read()

                                # Decrypt using XOR
                                decrypted_data = decrypt_binz(encrypted_data, encryption_key)

                                # Decompress (usually gzip)
                                decompressed_data = decompress_osgjs(decrypted_data)

                                # Check if it looks like valid JSON/OSGJS
                                if decompressed_data[:1] == b'{' or decompressed_data[:4] == b'osgj':
                                    osgjs_path = self.output_dir / 'model.osgjs'
                                    with open(osgjs_path, 'wb') as f:
                                        f.write(decompressed_data)
                                    files.append(str(osgjs_path))
                                    model_decrypted = True
                                else:
                                    # Save raw decrypted for debugging
                                    raw_path = self.output_dir / 'model.decrypted'
                                    with open(raw_path, 'wb') as f:
                                        f.write(decompressed_data)
                                    files.append(str(raw_path))
                            except Exception as e:
                                print(f"Decryption failed: {e}", file=sys.stderr)

                        # Try to decompress if .gz
                        elif model_ext == 'osgjs.gz':
                            try:
                                with gzip.open(model_path, 'rb') as f:
                                    osgjs_content = f.read()
                                osgjs_path = self.output_dir / 'model.osgjs'
                                with open(osgjs_path, 'wb') as f:
                                    f.write(osgjs_content)
                                files.append(str(osgjs_path))
                                model_decrypted = True
                            except:
                                pass

                # Save metadata
                metadata = {
                    'model_uid': model_uid,
                    'model_name': model_name,
                    'author': author,
                    'model_url': model_url,
                    'texture_count': downloaded_textures,
                    'texture_count_webgl': webgl_textures,
                    'texture_count_dump': dump_textures,
                    'obj_files': obj_files,
                    'model_downloaded': model_downloaded,
                    'model_decrypted': model_decrypted,
                    'encryption_key_found': len(encryption_key) > 0,
                    'intercepted_files': len(self.intercepted_bin_files),
                    'captured_geometry_info': captured_info,
                    'network_captures': self.captured_network_files if self.capture_network else [],
                    'source_url': f'https://sketchfab.com/3d-models/{model_uid}'
                }

                meta_path = self.output_dir / 'metadata.json'
                with open(meta_path, 'w') as f:
                    json.dump(metadata, f, indent=2)
                files.append(str(meta_path))

                return {
                    'success': True,
                    'model_uid': model_uid,
                    'model_name': model_name,
                    'author': author,
                    'files': files,
                    'texture_count': downloaded_textures,
                    'geometry_count': 1 if model_downloaded else 0,
                    'output_dir': str(self.output_dir)
                }

            except PlaywrightTimeout:
                await browser.close()
                return {
                    'success': False,
                    'error': f'Timeout after {self.timeout}s'
                }
            except Exception as e:
                await browser.close()
                return {
                    'success': False,
                    'error': str(e)
                }


async def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('model_uid')
    parser.add_argument('output_dir')
    parser.add_argument('--timeout', type=int, default=120)
    parser.add_argument('--headed', action='store_true')
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--capture-network', action='store_true')

    args = parser.parse_args()

    extractor = SketchfabExtractor(
        args.output_dir,
        args.timeout,
        args.headed,
        debug=args.debug,
        capture_network=args.capture_network
    )
    result = await extractor.run(args.model_uid)
    output_json(result)


if __name__ == '__main__':
    asyncio.run(main())
