from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse

import platform
import subprocess
import logging
import urllib.request
import os
import json
import uuid

import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler, EulerDiscreteScheduler


print(f"Is CUDA available: {torch.cuda.is_available()}")

app = FastAPI()

@app.get("/generate")
def generate_image(prompt, inference_steps, model):
    torch.cuda.empty_cache()
    print(f"Is CUDA available: {torch.cuda.is_available()}")

    pipeline = DiffusionPipeline.from_pretrained(str(model))
    
    #pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
    pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
    
    pipeline = pipeline.to("cuda")
    image = pipeline(prompt, num_inference_steps=int(inference_steps)).images[0]

    filename = str(uuid.uuid4()) + ".jpg"
    image.save(filename)
    
    print(filename)
    
    # Data to be written
    assertion = {
                "assertions": [
                    {
                    "label": "com.truepic.custom.ai",
                    "data": {
                        "model_name": model,
                        "model_version": "1.0",
                        "prompt": prompt
                    }
                    }
                ]
                }
     
    json_object = json.dumps(assertion, indent=4)
    with open("assertion.json", "w") as outfile:
        outfile.write(json_object)
    
    subprocess.check_output(['./truepic-sign', 'sign', filename, '--profile', 'demo', '--assertions', 'assertion.json', '--output', (os.getcwd() + '/static/' + filename)])

    return {"response": filename}


@app.get("/generate-picsum")
def generate_picsum(prompt):
    local_filename, headers = urllib.request.urlretrieve(('https://picsum.photos/id/' + prompt + '/800/800'))

    # Data to be written
    assertion = {
                "assertions": [
                    {
                    "label": "com.truepic.custom.ai",
                    "data": {
                        "model_name": "Picsum",
                        "model_version": "1.0",
                        "prompt": prompt
                    }
                    }
                ]
                }
    
    json_object = json.dumps(assertion, indent=4)
    with open("assertion.json", "w") as outfile:
        outfile.write(json_object)
    
    subprocess.check_output(['./truepic-sign', 'sign', local_filename, '--profile', 'demo', '--assertions', 'assertion.json', '--output', (os.getcwd() + '/static/output.jpg')])
    return {"response": "success"}

app.mount("/", StaticFiles(directory="static", html=True), name="static")

@app.get("/")
def index() -> FileResponse:
    return FileResponse(path="/app/static/index.html", media_type="text/html")