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 StableDiffusionPipeline, DPMSolverMultistepScheduler, EulerDiscreteScheduler

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 = StableDiffusionPipeline.from_pretrained(str(model), torch_dtype=torch.float16)
    
    #pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
    #another comment
    pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
    
    pipeline = pipeline.to("cuda")
    image = pipeline(prompt, num_inference_steps=int(inference_steps), height=512, width=512).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', '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")