import os
import time
import requests
import base64

from dotenv import load_dotenv

load_dotenv()
API_URL = os.environ["TRELLIS_URL"]


def generate_no_preview(image_base64: str):
    """Generate 3D model from a single base64-encoded image without previews.
    
    Args:
        image_base64: Base64 string of the image (without 'data:image/...' prefix)
    """
    try:
        # Set generation parameters
        params = {
            'image_base64': image_base64,
            'seed': 42,
            'ss_guidance_strength': 7.5,
            'ss_sampling_steps': 10,
            'slat_guidance_strength': 7.5,
            'slat_sampling_steps': 10,
            'mesh_simplify_ratio': 0.99,
            'texture_size': 1024,
            #'texture_baking_mode': 'opt',
            'texture_opt_total_steps': 1000,
            'output_format': 'glb'
        }
        
        # Start generation
        print("Starting generation...")
        response = requests.post(f"{API_URL}/generate_no_preview", data=params)
        response.raise_for_status()
        
        # Poll status until complete
        while True:
            status = requests.get(f"{API_URL}/status").json()
            print(f"Progress: {status['progress']}%")
            
            if status['status'] == 'COMPLETE':
                break
            elif status['status'] == 'FAILED':
                raise Exception(f"Generation failed: {status['message']}")
            
            time.sleep(1)
        
        # Download the model
        print("Downloading model...")
        response = requests.get(f"{API_URL}/download/model")
        response.raise_for_status()
        return response.content
        
        
    except Exception as e:
        print(f"Error: {str(e)}")
        return None
    
def image_to_3d_api(image_path, output_path):
    with open(image_path, 'rb') as image_file:
        image_data = image_file.read()

    base64_encoded = base64.b64encode(image_data).decode('utf-8')
    model_binary = generate_no_preview(base64_encoded)

    output_file = f"{output_path}.glb"
    with open(output_file, 'wb') as f:
        f.write(model_binary)

    return output_file