import sys
import os
import base64
import json
import urllib.request
import urllib.error
import argparse
import hashlib
import secrets

# Reconfigure stdout for UTF-8 on Windows command line
if sys.platform.startswith('win'):
    try:
        sys.stdout.reconfigure(encoding='utf-8')
        sys.stderr.reconfigure(encoding='utf-8')
    except Exception:
        pass

# ==========================================
# CONFIGURATION
# ==========================================
API_URL = "https://dr-fast-upload.pages.dev/api/upload"
MAX_SIZE_MB = 24
APP_CRYPTO_SECRET = "dr-fast-upload-secure-v1"
DRFU2_MAGIC = b"DRFU2\0\0\0"
DRFU2_CHUNK_SIZE = 1024 * 1024

def derive_key(secret: str) -> bytes:
    return hashlib.sha256(secret.encode('utf-8')).digest()

def encrypt_bytes(data: bytes, key: bytes) -> bytes:
    """
    Encrypts raw bytes using AES-256-GCM matching the Web Crypto format:
    12-byte IV + Ciphertext + 16-byte Auth Tag.
    Uses 'cryptography' library if installed, otherwise falls back to PyCryptodome.
    """
    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        aesgcm = AESGCM(key)
        base_nonce = secrets.token_bytes(8)
        output = bytearray(DRFU2_MAGIC + base_nonce)
        output += (DRFU2_CHUNK_SIZE).to_bytes(4, 'little')
        output += len(data).to_bytes(8, 'little')
        output += b'\0' * 4
        for index, offset in enumerate(range(0, len(data), DRFU2_CHUNK_SIZE)):
            chunk = data[offset:offset + DRFU2_CHUNK_SIZE]
            nonce = base_nonce + index.to_bytes(4, 'little')
            output += len(chunk).to_bytes(4, 'little')
            output += aesgcm.encrypt(nonce, chunk, None)
        if not data:
            output += (0).to_bytes(4, 'little')
            output += aesgcm.encrypt(base_nonce, b'', None)
        return bytes(output)
    except ImportError:
        try:
            from Crypto.Cipher import AES
            iv = secrets.token_bytes(12)
            cipher = AES.new(key, AES.MODE_GCM, nonce=iv)
            ciphertext, tag = cipher.encrypt_and_digest(data)
            return iv + ciphertext + tag
        except ImportError:
            print("❌ Error: Missing encryption library. Please install cryptography:")
            print("   pip install cryptography")
            sys.exit(1)

def upload_file(file_path):
    # 1. Check if the file exists
    if not os.path.isfile(file_path):
        print(f"❌ Error: The file '{file_path}' does not exist or is a directory.")
        sys.exit(1)

    # 2. Check the file size
    file_size_bytes = os.path.getsize(file_path)
    file_size_mb = file_size_bytes / (1024 * 1024)
    if file_size_mb > MAX_SIZE_MB:
        print(f"❌ Error: File is {file_size_mb:.2f}MB. Maximum allowed size is {MAX_SIZE_MB}MB.")
        sys.exit(1)

    filename = os.path.basename(file_path)
    print(f"⏳ Reading and encrypting '{filename}' ({file_size_mb:.2f} MB)...")

    # 3. Read raw file bytes and compute plaintext rawHash
    try:
        with open(file_path, "rb") as f:
            raw_bytes = f.read()
    except Exception as e:
        print(f"❌ Error reading file: {e}")
        sys.exit(1)

    raw_hash = hashlib.sha256(raw_bytes).hexdigest()

    # 4. Encrypt with AES-256-GCM
    key = derive_key(APP_CRYPTO_SECRET)
    encrypted_bytes = encrypt_bytes(raw_bytes, key)
    encoded_content = base64.b64encode(encrypted_bytes).decode('utf-8')

    # 5. Prepare JSON payload
    payload = json.dumps({
        "filename": filename,
        "content": encoded_content,
        "encrypted": True,
        "rawHash": raw_hash
    }).encode('utf-8')

    req = urllib.request.Request(API_URL, data=payload, headers={
        'Content-Type': 'application/json',
        'User-Agent': 'DrFastUpload-CLI/2.0'
    })

    # 5. Send to the server and handle the response
    try:
        with urllib.request.urlopen(req) as response:
            result = json.loads(response.read().decode('utf-8'))
            
            if result.get("success"):
                # Check if our backend deduplication caught it
                if result.get("cached"):
                    print("♻️  Already Exists! Instantly linked.")
                else:
                    print("✅ Upload Successful!")
                    
                print(f"🔗 File URL: {result.get('url')}")
            else:
                print(f"❌ Upload Failed: {result.get('error', 'Unknown error')}")

    except urllib.error.HTTPError as e:
        # Check if 409 Conflict contains existing file URL
        try:
            err_response = json.loads(e.read().decode('utf-8'))
            if err_response.get("cached") and err_response.get("url"):
                print("♻️  Already Exists! Instantly linked.")
                print(f"🔗 File URL: {err_response.get('url')}")
                return
            print(f"❌ Server Error: {err_response.get('error', e.reason)}")
        except Exception:
            print(f"❌ HTTP Error {e.code}: {e.reason}")
    except urllib.error.URLError as e:
        print(f"❌ Connection Error: Failed to reach the server. Please check your internet or the API_URL.")
    except Exception as e:
        print(f"❌ An unexpected error occurred: {str(e)}")

def main():
    # Setup the Help Menu
    parser = argparse.ArgumentParser(
        description="Dr Fast Upload CLI - Quickly host files from your terminal.",
        epilog="Examples:\n"
               "  python dr-fast-upload.py my_image.png\n"
               "  python dr-fast-upload.py \"C:\\Users\\Bob\\My Documents\\photo.jpg\"\n"
               "  python dr-fast-upload.py my file with spaces.pdf",
        formatter_class=argparse.RawTextHelpFormatter
    )
    
    # We use nargs='+' so that if a user forgets to use quotes around a file path 
    # with spaces (e.g., python script.py my photo.jpg), the script stitches it back together!
    parser.add_argument(
        "filepath", 
        nargs='+', 
        help="The path to the file you want to upload."
    )

    # If no arguments provided, print help
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()
    
    # Stitch the filepath back together in case they forgot quotes
    full_path = " ".join(args.filepath)
    
    upload_file(full_path)

if __name__ == "__main__":
    main()