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


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

    #model_id = "CompVis/stable-diffusion-v1-4" #stabilityai/stable-diffusion-2-1

    # Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here instead
    #pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    #pipe = StableDiffusionPipeline.from_pretrained(model_id)
    
    #pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    #pipe = pipe.to("cuda")
    
    pipeline = DiffusionPipeline.from_pretrained(str(model))
    pipeline = pipeline.to("cuda")
    #generator = torch.Generator("gpu").manual_seed(0)
    pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
    #image = pipeline(prompt, generator=generator).images[0]
    image = pipeline(prompt, num_inference_steps=int(inference_steps)).images[0]


    #prompt = "a photo of an astronaut riding a horse on mars"

    #image = pipe(prompt, num_inference_steps=5).images[0]
    #image = pipe(prompt).images[0]
        

    filename = str(uuid.uuid4()) + ".jpg"

    #print(f"after filename assignment")

    image.save(filename)
    
    print(filename)
    #print(f"after save")

    
    # 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")