import asyncio
import os
import sys
import httpx
from cyber_sdk.ali_oss import upload_localfile

async def upload_image(local_file_path: str, bucket_name: str = 'aigc-admin', bucket_path: str = 'template') -> str:
    """Uploads a local image to OSS and returns the CDN URL."""
    print(f"Uploading {local_file_path} to {bucket_name}/{bucket_path}...")
    # Use forward slashes so cyber_sdk's .split('/') can correctly extract filename on Windows
    safe_path = os.path.abspath(local_file_path).replace("\\", "/")
    result = await upload_localfile(
        file_path=safe_path, 
        bucket_path=bucket_path, 
        bucket_name=bucket_name
    )
    print("Upload SDK Response:", result)
    
    oss_object_key = result.get('oss_object_key')
    if oss_object_key:
        cdn_url = f"https://res.cybertogether.net/{oss_object_key}"
        return cdn_url
    return None

async def download_image(url: str, save_path: str):
    """Downloads an image from an HTTP link and saves it locally."""
    print(f"Downloading from {url} to {save_path}...")
    async with httpx.AsyncClient(timeout=30.0) as client:
        resp = await client.get(url)
        resp.raise_for_status()
        with open(save_path, 'wb') as f:
            f.write(resp.content)
    print(f"Download completed: {save_path}")

async def main():
    if len(sys.argv) < 2:
        print("Usage:")
        print("  Upload: python upload.py <file_path>")
        print("  Download: python upload.py download <url> <save_path>")
        
        # Self-test block if no param is given
        print("\n--- Running Self Test ---")
        test_file = 'img_1_gen.png'
        if not os.path.exists(test_file):
            print(f"Creating a dummy 1x1 PNG at {test_file} for testing.")
            with open(test_file, 'wb') as f:
                f.write(b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01\r\n-\xb4\x00\x00\x00\x00IEND\xaeB`\x82')
        
        url = await upload_image(test_file, 'aigc-admin', 'crawler/image')
        print(f"\nExtracted URL: {url}")
        
        if url:
            # Download it back
            download_path = "downloaded_dummy.png"
            await download_image(url, download_path)
            
    elif sys.argv[1] == 'download':
        if len(sys.argv) >= 4:
            await download_image(sys.argv[2], sys.argv[3])
        else:
            print("Error: Missing parameters for download.")
            print("Usage: python upload.py download <url> <save_path>")
    else:
        # Upload context
        file_path = sys.argv[1]
        url = await upload_image(file_path, 'aigc-admin', 'crawler/image')
        print(f"\nFinal CDN URL: {url}")

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