import os
import base64
import requests
import argparse
from pathlib import Path
from datetime import datetime
from dotenv import load_dotenv

load_dotenv()

ACCOUNT_ID = os.environ["CLOUDFLARE_ACCOUNT_ID"]
API_TOKEN = os.environ["CLOUDFLARE_API_TOKEN"]


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


def text_to_image(prompt, output_path):
    MODEL = "@cf/black-forest-labs/flux-1-schnell"
    URL = f"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/run/{MODEL}"

    payload = {
        "prompt": prompt,
    }

    headers = {
        "Authorization": f"Bearer {API_TOKEN}",
        "Content-Type": "application/json",
    }

    resp = requests.post(URL, json=payload, headers=headers, timeout=60)
    resp.raise_for_status()

    data = resp.json()
    b64 = data["result"]["image"]
    if not b64:
        raise RuntimeError(f"Unexpected response structure: {data}")

    img_bytes = base64.b64decode(b64)

    with open(output_path, "wb") as f:
        f.write(img_bytes)


def refine_text_prompt(prompt):
    MODEL = "@cf/meta/llama-3.2-3b-instruct"
    URL = f"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/run/{MODEL}"

    instructions = """
    User is talking about some object. Your task is to generate a short and concise description of it. Use only user's own words, keep it as short as possible.
    Example:
    User: 'Umm, okay, I would like a really cool sword, with for example a bright orange crossguard. And also it should be slightly curved.'
    You: 'a slightly curved sword with bright orange crossguard'
    """

    response = requests.post(URL,
    headers={"Authorization": f"Bearer {API_TOKEN}"},
    json={
        "messages": [
        {"role": "system", "content": instructions},
        {"role": "user", "content": prompt}
        ]
    }
    )
    data = response.json()
    return data["result"]["response"]

def image_to_3d(image_path, output_path):
    import subprocess

    # TODO: make this configurable
    MODEL_FOLDER = r"D:\henrisel\AI-models\stable-fast-3d"

    venv_python = MODEL_FOLDER + r"\.venv\Scripts\python.exe"
    script_path = MODEL_FOLDER + r"\run.py"

    args = [image_path, "--output-dir", output_path]
    command = [venv_python, script_path] + args

    try:
        # Run the subprocess
        result = subprocess.run(command, capture_output=True, text=True)

        # Print output and errors
        print("STDOUT:\n", result.stdout)
        print("STDERR:\n", result.stderr)
        print("Return Code:", result.returncode)

    except Exception as e:
        print(f"Error occurred: {e}")

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()

    user_prompt = args.prompt
    print(f"User prompt: {user_prompt}")
    refined_prompt = refine_text_prompt(user_prompt)
    print(f"Refined prompt: {refined_prompt}")
    timestamp = get_timestamp()
    image_path = Path.cwd() / "images" / f"{timestamp}.jpg"
    text_to_image(refined_prompt, image_path)
    print(f"Generated image file: {image_path}")
    model_path = Path.cwd() / "models" / timestamp
    image_to_3d(image_path, model_path)
    print(f"Generated 3D model file: {model_path}")


if __name__ == "__main__":
    main()