import argparse
import asyncio
import logging
import time

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

from generate_image_local import text_to_image_invoke_ai
from generate_model_local import image_to_3d_api

load_dotenv()


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

def setup_logger(base_folder, timestamp):
    log_dir = base_folder / Path("logs")
    log_dir.mkdir(parents=True, exist_ok=True)
    logging.basicConfig(
        filename=log_dir / f"{timestamp}.log",
        level=logging.INFO,
        force=True
    )


async 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}")
    image_generation_prompt = input_prompt + ", single object, fully visible, plain neutral background, high-resolution, sharp details, soft studio lighting, front and side visible, realistic style, true-to-scale"
    
    pipeline_folder = Path(__file__).resolve().parent
    timestamp = get_timestamp()
    setup_logger(pipeline_folder, timestamp)
    time_checkpoint = time.time()

    image_path = pipeline_folder / "images" / f"{timestamp}.jpg"
    # TODO: use Invoke AI or Cloudflare, depending on env var
    #text_to_image_cloudflare(image_generation_prompt, image_path)
    await text_to_image_invoke_ai(image_generation_prompt, image_path)
    
    image_generation_time = time.time() - time_checkpoint
    time_checkpoint = time.time()
    logging.info(f"Image generation time: {round(image_generation_time, 1)} s")
    print(f"Generated image file: {image_path}")

    model_path = pipeline_folder / "models" / timestamp
    model_file = image_to_3d_api(image_path, model_path)

    model_generation_time = time.time() - time_checkpoint
    logging.info(f"Model generation time: {round(model_generation_time, 1)} s")
    print(f"Generated 3D model file: {model_file}")


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