#!/usr/bin/env python3
"""
LaoZhang.ai API client for Veo 3.1 video generation.

OpenAI-compatible API with image-to-video support (frame interpolation).
Docs: https://docs.laozhang.ai/
"""

import base64
import json
import os
import re
import time
import requests
from pathlib import Path
from typing import Optional
import ssl
from requests.adapters import HTTPAdapter
from urllib3.util.ssl_ import create_urllib3_context

# Default proxy settings for bypassing SSL interception
# Set PROXY_HOST=none to disable proxy
DEFAULT_PROXY_HOST = os.environ.get("PROXY_HOST", "none")
DEFAULT_PROXY_PORT = int(os.environ.get("PROXY_PORT", "1080"))


def extract_url_from_markdown(text: str) -> Optional[str]:
    """Extract URL from markdown link format [text](url) or plain URL."""
    # Try markdown link format first
    match = re.search(r'\[.*?\]\((https?://[^\)]+)\)', text)
    if match:
        return match.group(1)

    # Try plain URL
    match = re.search(r'(https?://\S+\.mp4)', text)
    if match:
        return match.group(1)

    return None


def extract_base64_image_from_markdown(text: str) -> Optional[tuple]:
    """
    Extract base64 image data from markdown format.

    Returns:
        Tuple of (image_bytes, mime_type) or None if not found
    """
    # Match markdown image with base64 data URL: ![...](data:image/png;base64,...)
    match = re.search(r'!\[.*?\]\((data:image/(\w+);base64,([A-Za-z0-9+/=]+))\)', text)
    if match:
        mime_type = f"image/{match.group(2)}"
        base64_data = match.group(3)
        image_bytes = base64.b64decode(base64_data)
        return image_bytes, mime_type

    return None

BASE_URL = "https://api.laozhang.ai/v1"
BASE_URL_NATIVE = "https://api.laozhang.ai/v1beta"


class RetryableVeo3Error(Exception):
    """Error that can be retried (e.g., PUBLIC_ERROR_IP_INPUT_IMAGE)."""
    pass


class TLSAdapter(HTTPAdapter):
    """
    HTTP Adapter with browser-like TLS settings.
    """
    def init_poolmanager(self, *args, **kwargs):
        # Create SSL context similar to browsers
        ctx = create_urllib3_context()
        ctx.minimum_version = ssl.TLSVersion.TLSv1_2
        ctx.set_ciphers('DEFAULT:@SECLEVEL=1')  # More permissive cipher suite
        kwargs['ssl_context'] = ctx
        return super().init_poolmanager(*args, **kwargs)


def _extract_laozhang_error(response, last_error=None) -> str:
    """Extract human-readable error from LaoZhang API response."""
    if response is None:
        return str(last_error) if last_error else "no response"
    try:
        data = response.json()
        err = data.get("error", {})
        msg = err.get("message", "")
        localized = err.get("localized_message", "")
        code = err.get("code", "")
        parts = [p for p in [msg, localized, code] if p]
        if parts:
            return " | ".join(parts)
    except Exception:
        pass
    return response.text[:500]


def get_laozhang_session() -> requests.Session:
    """Get a requests session with browser-like TLS settings."""
    session = requests.Session()
    adapter = TLSAdapter()
    session.mount("https://", adapter)
    return session


def get_proxied_laozhang_session(
    proxy_host: str = DEFAULT_PROXY_HOST,
    proxy_port: int = DEFAULT_PROXY_PORT
) -> requests.Session:
    """
    Get a requests session that routes through SOCKS5 proxy.

    Used for /v1beta endpoint which requires bypassing SSL interception.
    If proxy_host is empty or "none", proxy is disabled.

    Note: TLSAdapter is NOT used with proxy - it breaks SSL handshake.
    """
    session = requests.Session()

    # Configure SOCKS5 proxy if specified
    # Use socks5h:// for remote DNS resolution (required for proper SSL)
    if proxy_host and proxy_host.lower() != "none":
        proxy_url = f"socks5h://{proxy_host}:{proxy_port}"
        session.proxies = {
            "http": proxy_url,
            "https": proxy_url
        }
    else:
        # Only use TLSAdapter when NOT using proxy
        adapter = TLSAdapter()
        session.mount("https://", adapter)

    return session


def image_to_base64_url(image_path: str) -> str:
    """Convert image file to base64 data URL."""
    path = Path(image_path)
    with open(path, 'rb') as f:
        image_data = f.read()

    suffix = path.suffix.lower()
    if suffix in ['.jpg', '.jpeg']:
        mime_type = 'image/jpeg'
    elif suffix == '.png':
        mime_type = 'image/png'
    elif suffix == '.webp':
        mime_type = 'image/webp'
    else:
        mime_type = 'image/png'

    image_base64 = base64.b64encode(image_data).decode('utf-8')
    return f"data:{mime_type};base64,{image_base64}"


def call_nanobanana_laozhang(
    image_data: bytes,
    prompt: str,
    api_key: str,
    mime_type: str = "image/webp",
    aspect_ratio: str = "16:9",
    image_size: str = "1K",
    proxy_host: str = DEFAULT_PROXY_HOST,
    proxy_port: int = DEFAULT_PROXY_PORT,
    reference_image_data: bytes = None,
    reference_mime_type: str = None
) -> tuple:
    """
    Call NanoBanana Pro (Gemini 3 Pro image-to-image) via LaoZhang.ai.

    Uses Google native format to support custom aspect ratios.
    Routes through SOCKS5 proxy to bypass local SSL interception.

    Args:
        image_data: Raw image bytes (main product image)
        prompt: Editing/processing instructions
        api_key: LaoZhang API key
        mime_type: Input image MIME type
        aspect_ratio: Output aspect ratio ("16:9", "9:16", "1:1", "4:3", "3:4")
        image_size: Output size ("1K", "2K", "4K")
        proxy_host: SOCKS5 proxy host (default: proxy-host)
        proxy_port: SOCKS5 proxy port (default: 1080)
        reference_image_data: Optional reference image for scene consistency
        reference_mime_type: MIME type of reference image

    Returns:
        Tuple of (image_bytes, mime_type)
    """
    session = get_proxied_laozhang_session(proxy_host, proxy_port)

    image_base64 = base64.b64encode(image_data).decode('utf-8')

    # Build parts list - main image first
    parts = [
        {"inline_data": {"mime_type": mime_type, "data": image_base64}},
    ]

    # Add reference image if provided (for scene consistency)
    if reference_image_data and reference_mime_type:
        ref_base64 = base64.b64encode(reference_image_data).decode('utf-8')
        parts.append({"inline_data": {"mime_type": reference_mime_type, "data": ref_base64}})
        print(f"  Using reference image for scene consistency")

    # Add prompt last
    parts.append({"text": prompt})

    # Google native format for aspect ratio support
    payload = {
        "contents": [{
            "parts": parts
        }],
        "generationConfig": {
            "responseModalities": ["IMAGE"],
            "imageConfig": {
                "aspectRatio": aspect_ratio,
                "imageSize": image_size
            }
        }
    }

    print(f"  Calling NanoBanana Pro via LaoZhang.ai (proxy: {proxy_host}:{proxy_port})...")
    print(f"  Aspect ratio: {aspect_ratio}, Size: {image_size}")

    # Retry logic for intermittent SSL/502/503 errors
    max_retries = 3
    last_error = None
    response = None
    for attempt in range(max_retries):
        try:
            response = session.post(
                f"{BASE_URL_NATIVE}/models/gemini-3-pro-image-preview:generateContent",
                headers={
                    "x-goog-api-key": api_key,
                    "Content-Type": "application/json"
                },
                json=payload,
                timeout=180
            )
            # Check for retryable status codes
            if response.status_code in (429, 502, 503, 504):
                last_error = f"HTTP {response.status_code}"
                if attempt < max_retries - 1:
                    wait = 10 if response.status_code == 429 else 5
                    print(f"  Retry {attempt + 1}/{max_retries} after {last_error} (wait {wait}s)")
                    time.sleep(wait)
                    continue
            break  # Success or non-retryable error
        except Exception as e:
            last_error = str(e)
            if attempt < max_retries - 1:
                print(f"  Retry {attempt + 1}/{max_retries} after error: {e}")
                time.sleep(5)
                continue
    else:
        # All retries exhausted via exceptions
        raise RuntimeError(f"NanoBanana failed after {max_retries} retries: {last_error}")

    if response is None or response.status_code != 200:
        status = response.status_code if response else "no response"
        error_detail = _extract_laozhang_error(response, last_error)
        raise RuntimeError(f"NanoBanana LaoZhang API error: HTTP {status} - {error_detail}")

    data = response.json()

    # Extract image from response
    try:
        parts = data["candidates"][0]["content"]["parts"]
        if not parts:
            raise KeyError("parts is null/empty (likely content filtered)")
        for part in parts:
            if "inlineData" in part:
                image_bytes = base64.b64decode(part["inlineData"]["data"])
                returned_mime = part["inlineData"].get("mimeType", "image/png")
                print(f"  NanoBanana returned image: {len(image_bytes) // 1024}KB, {returned_mime}")
                return image_bytes, returned_mime
        raise KeyError("No image in response")
    except (KeyError, IndexError, TypeError) as e:
        # Check if there's an error message
        error_msg = data.get("error", {}).get("message", str(data)[:500])
        raise RuntimeError(f"Failed to extract image from NanoBanana response: {e}\n{error_msg}")


def call_seedream_laozhang(
    prompt: str,
    api_key: str,
    aspect_ratio: str = "3:4",
    image_url: str = None,
) -> tuple:
    """
    Call SeeDream 4.5 (ByteDance) via LaoZhang.ai for image generation.

    Used as fallback when NanoBanana fails with IMAGE_SAFETY filters.

    Args:
        prompt: Image generation prompt
        api_key: LaoZhang API key
        aspect_ratio: Output aspect ratio ("3:4", "9:16", "16:9")
        image_url: Optional reference image URL for img2img

    Returns:
        Tuple of (image_bytes, mime_type)
    """
    # Map aspect ratio to width/height
    aspect_map = {
        "3:4": (768, 1024),
        "9:16": (720, 1280),
        "16:9": (1280, 720),
        "4:3": (1024, 768),
        "1:1": (1024, 1024),
    }
    w, h = aspect_map.get(aspect_ratio, (768, 1024))

    # Add portrait hint for vertical images
    if h > w:
        prompt = f"Generate a VERTICAL PORTRAIT image (taller than wide). {prompt}"

    payload = {
        "model": "seedream-4-5-251128",
        "prompt": prompt,
        "response_format": "url",
        "size": "2K",
        "width": w,
        "height": h,
        "watermark": False,
    }
    if image_url:
        payload["image"] = image_url

    session = get_laozhang_session()

    print(f"  Calling SeeDream 4.5 via LaoZhang.ai...")
    print(f"  Size: {w}x{h}, img2img: {bool(image_url)}")

    max_retries = 2
    last_error = None
    response = None

    for attempt in range(max_retries):
        try:
            response = session.post(
                f"{BASE_URL}/images/generations",
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json"
                },
                json=payload,
                timeout=180
            )
            if response.status_code in (429, 502, 503):
                last_error = f"HTTP {response.status_code}"
                if attempt < max_retries - 1:
                    print(f"  Retry {attempt + 1}/{max_retries} after {last_error}")
                    time.sleep(5)
                    continue
            break
        except Exception as e:
            last_error = str(e)
            if attempt < max_retries - 1:
                print(f"  Retry {attempt + 1}/{max_retries} after error: {e}")
                time.sleep(5)
                continue
    else:
        raise RuntimeError(f"SeeDream failed after {max_retries} retries: {last_error}")

    if response is None or response.status_code != 200:
        status = response.status_code if response else "no response"
        error_detail = _extract_laozhang_error(response, last_error)
        raise RuntimeError(f"SeeDream API error: HTTP {status} - {error_detail}")

    data = response.json()

    # Extract image URL from response
    try:
        image_url_result = data["data"][0]["url"]
    except (KeyError, IndexError) as e:
        raise RuntimeError(f"Failed to extract image URL from SeeDream response: {e}\n{str(data)[:500]}")

    # Download image from URL
    print(f"  Downloading SeeDream image...")
    img_response = session.get(image_url_result, timeout=60)
    if img_response.status_code != 200:
        raise RuntimeError(f"Failed to download SeeDream image: HTTP {img_response.status_code}")

    image_bytes = img_response.content
    print(f"  SeeDream returned image: {len(image_bytes) // 1024}KB")
    return image_bytes, "image/jpeg"


def call_veo3_laozhang(
    prompt: str,
    api_key: str,
    first_frame_path: Optional[str] = None,
    last_frame_path: Optional[str] = None,
    model: str = "veo-3.1-fl",
    n: int = 1,
    duration: Optional[int] = None,
    proxy_host: str = DEFAULT_PROXY_HOST,
    proxy_port: int = DEFAULT_PROXY_PORT
) -> dict:
    """
    Generate video using LaoZhang.ai Veo 3.1 API.

    Args:
        prompt: Video generation prompt
        api_key: LaoZhang API key
        first_frame_path: Path to first frame image (optional)
        last_frame_path: Path to last frame image (optional, for interpolation)
        model: Model to use:
            - veo-3.1-landscape-fast-fl: Landscape, fast, with images ($0.15)
            - veo-3.1-fast-fl: Fast with images ($0.15)
            - veo-3.1-fl: Standard with images ($0.25)
            - veo-3.1-fast: Fast text-only ($0.15)
            - veo-3.1: Standard text-only ($0.25)
        n: Number of videos to generate (1-4)
        duration: Video duration in seconds (default: 8)

    Returns:
        Dict with video URL(s)
    """
    # Build content array
    content = [{"type": "text", "text": prompt}]

    # Add first frame if provided
    if first_frame_path and Path(first_frame_path).exists():
        content.append({
            "type": "image_url",
            "image_url": {"url": image_to_base64_url(first_frame_path)}
        })
        print(f"  Added first frame: {Path(first_frame_path).name}")

    # Add last frame if provided (for interpolation)
    if last_frame_path and Path(last_frame_path).exists():
        content.append({
            "type": "image_url",
            "image_url": {"url": image_to_base64_url(last_frame_path)}
        })
        print(f"  Added last frame: {Path(last_frame_path).name}")

    payload = {
        "model": model,
        "messages": [{
            "role": "user",
            "content": content
        }],
        "stream": True,
        "n": n
    }

    # Add duration if specified (5-8 seconds supported)
    if duration:
        payload["duration"] = duration

    print(f"  Model: {model}")
    if duration:
        print(f"  Duration: {duration}s")
    print(f"  Generating {n} video(s)...")

    # Call Veo3 API (retries with image regeneration handled in full_pipeline)
    return _call_veo3_laozhang_once(payload, proxy_host, proxy_port, api_key)


def _call_veo3_laozhang_once(
    payload: dict,
    proxy_host: str,
    proxy_port: int,
    api_key: str
) -> dict:
    """Single attempt to call Veo3 API."""
    # Make streaming request via proxy (bypasses local DNS/SSL issues)
    session = get_proxied_laozhang_session(proxy_host, proxy_port)
    response = session.post(
        f"{BASE_URL}/chat/completions",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        },
        json=payload,
        stream=True,
        timeout=600  # 10 minutes for video generation
    )

    if response.status_code != 200:
        error_detail = _extract_laozhang_error(response)
        raise RuntimeError(f"LaoZhang API error: HTTP {response.status_code} - {error_detail}")

    # Process streaming response
    video_urls = []
    last_progress = ""

    for line in response.iter_lines():
        if not line:
            continue

        line_text = line.decode('utf-8')
        if not line_text.startswith('data: '):
            continue

        data_str = line_text[6:]  # Remove 'data: ' prefix
        if data_str == '[DONE]':
            break

        try:
            data = json.loads(data_str)
            # Write debug to file
            with open("/tmp/laozhang_debug.log", "a") as f:
                f.write(f"{json.dumps(data)}\n")

            # Check for completion
            choices = data.get('choices', [])
            for choice in choices:
                # Check for video URL in message
                message = choice.get('message', {})
                if message:
                    content = message.get('content', '')
                    if content and ('http' in content or content.endswith('.mp4')):
                        video_urls.append(content)
                        print(f"  Video ready: {content[:80]}...")

                # Check delta for streaming progress
                delta = choice.get('delta', {})
                if delta:
                    content = delta.get('content', '')
                    if content:
                        # Check for errors in stream
                        if '错误' in content or 'ERROR' in content or '❌' in content:
                            # Check for retryable errors
                            if 'PUBLIC_ERROR_IP_INPUT_IMAGE' in content:
                                raise RetryableVeo3Error(f"Veo3 generation error: {content.strip()}")
                            raise RuntimeError(f"Veo3 generation error: {content.strip()}")
                        # Print progress updates (but not too often)
                        if content != last_progress and len(content) < 200:
                            if 'progress' in content.lower() or '%' in content or '进度' in content:
                                print(f"  {content.strip()}")
                            last_progress = content
                        # Check for URL in delta
                        if 'http' in content or content.endswith('.mp4'):
                            video_urls.append(content)

                # Check finish_reason
                if choice.get('finish_reason') == 'stop':
                    # Final message might contain URL
                    pass

        except json.JSONDecodeError:
            continue

    if not video_urls:
        raise RuntimeError("No video URL received from API")

    # Extract actual URLs from markdown format if needed
    extracted_urls = []
    for url_text in video_urls:
        extracted = extract_url_from_markdown(url_text)
        if extracted:
            extracted_urls.append(extracted)
        elif url_text.startswith('http'):
            extracted_urls.append(url_text)

    if not extracted_urls:
        raise RuntimeError(f"Could not extract video URL from: {video_urls}")

    return {
        "uri": extracted_urls[0],
        "urls": extracted_urls,
        "state": "COMPLETE"
    }


def transcribe_audio_laozhang(
    audio_path: str,
    api_key: str,
    model: str = "whisper-1",
    language: str = None,
    proxy_host: str = DEFAULT_PROXY_HOST,
    proxy_port: int = DEFAULT_PROXY_PORT
) -> str:
    """
    Transcribe audio file using LaoZhang.ai Whisper API.

    Args:
        audio_path: Path to audio file (mp3, ogg, wav, etc.)
        api_key: LaoZhang API key
        model: Model to use (whisper-1, gpt-4o-transcribe)
        language: Optional language hint (e.g., "ru", "en")
        proxy_host: SOCKS5 proxy host
        proxy_port: SOCKS5 proxy port

    Returns:
        Transcribed text
    """
    session = get_proxied_laozhang_session(proxy_host, proxy_port)

    max_retries = 3
    last_error = None

    for attempt in range(max_retries):
        try:
            with open(audio_path, 'rb') as audio_file:
                files = {
                    'file': (Path(audio_path).name, audio_file, 'audio/ogg'),
                    'model': (None, model),
                }
                if language:
                    files['language'] = (None, language)

                response = session.post(
                    f"{BASE_URL}/audio/transcriptions",
                    headers={
                        "Authorization": f"Bearer {api_key}",
                    },
                    files=files,
                    timeout=60
                )

            if response.status_code in (502, 503, 504):
                last_error = f"HTTP {response.status_code}"
                if attempt < max_retries - 1:
                    time.sleep(3)
                    continue
            break

        except Exception as e:
            last_error = e
            if attempt < max_retries - 1:
                time.sleep(3)
                continue
            raise RuntimeError(f"Transcription failed after {max_retries} retries: {e}")

    if response.status_code != 200:
        raise RuntimeError(f"Transcription error: {response.status_code} - {response.text[:500]}")

    data = response.json()
    return data.get("text", "")


def transcribe_audio_with_timestamps(
    audio_path: str,
    api_key: str,
    model: str = "whisper-1",
    language: str = None,
    proxy_host: str = DEFAULT_PROXY_HOST,
    proxy_port: int = DEFAULT_PROXY_PORT
) -> dict:
    """
    Transcribe audio file with timestamps using LaoZhang.ai Whisper API.

    Args:
        audio_path: Path to audio file (mp3, ogg, wav, etc.)
        api_key: LaoZhang API key
        model: Model to use (whisper-1, gpt-4o-transcribe)
        language: Optional language hint (e.g., "ru", "en")
        proxy_host: SOCKS5 proxy host
        proxy_port: SOCKS5 proxy port

    Returns:
        Dict with 'text' (full transcription) and 'segments' (list of timestamped segments).
        Each segment has: {'start': float, 'end': float, 'text': str}
    """
    session = get_proxied_laozhang_session(proxy_host, proxy_port)

    max_retries = 3
    last_error = None

    for attempt in range(max_retries):
        try:
            with open(audio_path, 'rb') as audio_file:
                files = {
                    'file': (Path(audio_path).name, audio_file, 'audio/ogg'),
                    'model': (None, model),
                    'response_format': (None, 'verbose_json'),
                }
                if language:
                    files['language'] = (None, language)

                response = session.post(
                    f"{BASE_URL}/audio/transcriptions",
                    headers={
                        "Authorization": f"Bearer {api_key}",
                    },
                    files=files,
                    timeout=60
                )

            if response.status_code in (502, 503, 504):
                last_error = f"HTTP {response.status_code}"
                if attempt < max_retries - 1:
                    time.sleep(3)
                    continue
            break

        except Exception as e:
            last_error = e
            if attempt < max_retries - 1:
                time.sleep(3)
                continue
            raise RuntimeError(f"Transcription failed after {max_retries} retries: {e}")

    if response.status_code != 200:
        raise RuntimeError(f"Transcription error: {response.status_code} - {response.text[:500]}")

    data = response.json()

    # Extract segments with timestamps
    segments = []
    for seg in data.get("segments", []):
        segments.append({
            "start": seg.get("start", 0),
            "end": seg.get("end", 0),
            "text": seg.get("text", "").strip()
        })

    return {
        "text": data.get("text", ""),
        "segments": segments,
        "language": data.get("language", ""),
        "duration": data.get("duration", 0)
    }


def download_video_laozhang(video_url: str, output_path: str, max_retries: int = 3) -> str:
    """Download video from URL with proxy support and validation.

    Uses LAOZHANG_CDN_PROXY env var for SOCKS5 proxy if set.
    Validates downloaded file size matches Content-Length.
    """
    cdn_proxy = os.environ.get("LAOZHANG_CDN_PROXY", "")
    proxies = {"http": cdn_proxy, "https": cdn_proxy} if cdn_proxy else None

    last_error = None
    for attempt in range(1, max_retries + 1):
        try:
            response = requests.get(
                video_url, stream=True, timeout=(15, 120),
                proxies=proxies
            )

            if response.status_code != 200:
                raise RuntimeError(f"Download error: HTTP {response.status_code}")

            expected_size = int(response.headers.get("content-length", 0))

            with open(output_path, 'wb') as f:
                downloaded = 0
                for chunk in response.iter_content(chunk_size=65536):
                    f.write(chunk)
                    downloaded += len(chunk)

            # Validate file size
            if expected_size > 0 and downloaded != expected_size:
                raise RuntimeError(
                    f"Incomplete download: got {downloaded} bytes, expected {expected_size}"
                )

            # Sanity check: video should be at least 100KB
            if downloaded < 100_000:
                raise RuntimeError(
                    f"Downloaded file too small: {downloaded} bytes (likely corrupted)"
                )

            return output_path

        except Exception as e:
            last_error = e
            if attempt < max_retries:
                wait = attempt * 5
                print(f"  Download attempt {attempt}/{max_retries} failed: {e}, retrying in {wait}s...")
                import time
                time.sleep(wait)
            # Clean up partial file
            if os.path.exists(output_path):
                os.remove(output_path)

    raise RuntimeError(f"Download failed after {max_retries} attempts: {last_error}")


def test_laozhang(api_key: str) -> dict:
    """Test LaoZhang API connection."""
    try:
        session = get_laozhang_session()
        # Simple text generation test (cheaper)
        response = session.post(
            f"{BASE_URL}/chat/completions",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            },
            json={
                "model": "gpt-3.5-turbo",
                "messages": [{"role": "user", "content": "Say 'OK'"}],
                "max_tokens": 5
            },
            timeout=30
        )

        if response.status_code == 200:
            return {"status": "ok", "response": response.json()}
        else:
            return {"status": "error", "code": response.status_code, "text": response.text[:200]}
    except Exception as e:
        return {"status": "error", "message": str(e)}


if __name__ == "__main__":
    import sys
    import os
    from dotenv import load_dotenv

    load_dotenv(Path(__file__).parent.parent / ".env")

    api_key = os.environ.get("LAOZHANG_KEY")
    if not api_key:
        print("Error: LAOZHANG_KEY not set in .env")
        sys.exit(1)

    if len(sys.argv) > 1:
        cmd = sys.argv[1]

        if cmd == "test":
            print("Testing LaoZhang API (OpenAI endpoint)...")
            result = test_laozhang(api_key)
            print(f"Result: {result}")

        elif cmd == "test-nanobanana":
            print("Testing NanoBanana via proxy...")
            # Simple text-to-image test (no input image)
            try:
                session = get_proxied_laozhang_session()
                response = session.post(
                    f"{BASE_URL_NATIVE}/models/gemini-3-pro-image-preview:generateContent",
                    headers={
                        "x-goog-api-key": api_key,
                        "Content-Type": "application/json"
                    },
                    json={
                        "contents": [{"parts": [{"text": "A simple red circle on white background, 16:9 aspect ratio"}]}],
                        "generationConfig": {
                            "responseModalities": ["IMAGE"],
                            "imageConfig": {"aspectRatio": "16:9"}
                        }
                    },
                    timeout=120
                )
                print(f"Status: {response.status_code}")
                if response.status_code == 200:
                    data = response.json()
                    parts = data.get("candidates", [{}])[0].get("content", {}).get("parts", [])
                    for part in parts:
                        if "inlineData" in part:
                            size = len(part["inlineData"]["data"]) * 3 // 4  # base64 to bytes
                            mime = part["inlineData"].get("mimeType", "unknown")
                            print(f"Success! Image: {size // 1024}KB, {mime}")
                            break
                    else:
                        print(f"Response (no image): {str(data)[:500]}")
                else:
                    print(f"Error: {response.text[:500]}")
            except Exception as e:
                print(f"Error: {e}")

        else:
            print("Usage: python laozhang_proxy.py [test|test-nanobanana]")
    else:
        print("Usage: python laozhang_proxy.py [test|test-nanobanana]")
