import os
import argparse

from pathlib import Path
from datetime import datetime
from dotenv import load_dotenv

from cloudflare_api import text_to_image, refine_text_prompt
from generate_model_local import image_to_3d_api, image_to_3d_subprocess

load_dotenv()

PIPELINE_FOLDER = os.environ["PIPELINE_FOLDER"]


def get_timestamp():
    return datetime.now().strftime("%Y-%m-%d-%H-%M-%S")


def main():
    parser = argparse.ArgumentParser(description="Text to 3D model pipeline")
    parser.add_argument("--prompt", type=str, required=True, help="User text prompt")
    args = parser.parse_args()

    input_prompt = args.prompt
    print(f"Input prompt: {input_prompt}")

    refine_prompt = os.environ["REFINE_PROMPT"] == "1"
    if refine_prompt:
        image_generation_prompt = refine_text_prompt(input_prompt)
        print(f"Refined prompt: {image_generation_prompt}")
    else:
        image_generation_prompt = input_prompt
    
    timestamp = get_timestamp()
    pipeline_folder = Path(PIPELINE_FOLDER)
    image_path = pipeline_folder / "images" / f"{timestamp}.jpg"
    text_to_image(image_generation_prompt, image_path)
    print(f"Generated image file: {image_path}")
    model_path = pipeline_folder / "models" / timestamp
    image_to_3d_api(image_path, model_path)
    #model_file_path = model_path / "0" / "mesh.glb"
    print(f"Generated 3D model file: {model_path}")


if __name__ == "__main__":
    main()