from fastapi import FastAPI, Request, UploadFile, HTTPException, status
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import aiofiles

import subprocess
import os
import json
import uuid
import logging

import torch
from diffusers import (
    StableDiffusionPipeline,
    DPMSolverMultistepScheduler,
    EulerDiscreteScheduler,
)

app = FastAPI()


@app.get("/generate")
def generate_image(prompt, model):
    torch.cuda.empty_cache()

    modelArray = model.split(",")
    modelName = modelArray[0]
    modelVersion = modelArray[1]

    pipeline = StableDiffusionPipeline.from_pretrained(
        str(modelName), torch_dtype=torch.float16
    )
    pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
    pipeline = pipeline.to("cuda")

    image = pipeline(prompt, num_inference_steps=50, height=512, width=512).images[0]

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

    assertion = {
        "assertions": [
            {
                "label": "com.truepic.custom.ai",
                "data": {
                    "model_name": modelName,
                    "model_version": modelVersion,
                    "prompt": prompt,
                },
            }
        ]
    }

    json_object = json.dumps(assertion)

    subprocess.check_output(
        [
            "./truepic",
            "sign",
            filename,
            "--assertions-inline",
            json_object,
            "--output",
            (os.getcwd() + "/static/" + filename),
        ]
    )

    return {"response": filename}

@app.post("/verify")
def verify_image(file: UploadFile):
    logging.warning("in verify")
    logging.warning(file.filename)
    
    logging.warning('form')
    logging.warning(form)

    logging.warning('fileitem')
    logging.warning(fileitem)

    # check if the file has been uploaded
    if fileitem.filename:
        # strip the leading path from the file name
        fn = os.path.basename(fileitem.filename)
        
        # open read and write the file into the server
        open(fn, 'wb').write(fileitem.file.read())
        
    return {"response": fileitem.filename}

@app.post('/upload')
async def upload(fileUpload: UploadFile):
    try:
        contents = await fileUpload.read()
        async with aiofiles.open(fileUpload.filename, 'wb') as f:
            await f.write(contents)
    except Exception:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail='There was an error uploading the file',
        )
    finally:
        await fileUpload.close()

    return {'message': f'Successfuly uploaded {fileUpload.filename}'}

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