from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import shutil
import os
import json
import gradio as gr

app = FastAPI(docs_url=None, redoc_url=None)

# Create folders to store uploaded images
UPLOAD_FOLDER = "uploads"
IMAGE_FOLDER = os.path.join(UPLOAD_FOLDER, "images")
VIDEO_FOLDER = os.path.join(UPLOAD_FOLDER, "videos")

for folder in [UPLOAD_FOLDER, IMAGE_FOLDER, VIDEO_FOLDER]:
    os.makedirs(folder, exist_ok=True)

# Create a dictionary to store available images
available_images = {}


class ImageInfo(BaseModel):
    url: str
    name: str


@app.post("/upload/")
async def upload_image(url: str, file: UploadFile = File(...)):
    allowed_formats = [".png", ".jpg"]
    file_format = os.path.splitext(file.filename)[1]

    if file_format not in allowed_formats:
        raise HTTPException(status_code=400, detail="Unsupported file format")

    # Save the uploaded file
    with open(os.path.join(UPLOAD_FOLDER, file.filename), "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    # Determine the folder based on the file format
    folder = "images" if file_format in [".png", ".jpg"] else "videos"

    # Move the file to the appropriate folder
    os.rename(
        os.path.join(UPLOAD_FOLDER, file.filename),
        os.path.join(UPLOAD_FOLDER, folder, file.filename)
    )

    # Update available images dictionary
    available_images[file.filename] = ImageInfo(url=url, name=file.filename)

    return JSONResponse(content={"message": "File uploaded successfully"})


@app.get("/get-all-images/")
async def get_all_images(password: str, from_folder: str):
    if password != "Kastg@123":
        raise HTTPException(status_code=401, detail="Unauthorized")

    folder_path = os.path.join(UPLOAD_FOLDER, from_folder)
    if not os.path.exists(folder_path):
        raise HTTPException(status_code=404, detail="Folder not found")

    images = os.listdir(folder_path)
    return JSONResponse(content={"images": images})


@app.delete("/delete-image/")
async def delete_image(image_name: str, password: str):
    if password != "Kastg@123":
        raise HTTPException(status_code=401, detail="Unauthorized")

    image_path = os.path.join(UPLOAD_FOLDER, image_name)
    if not os.path.exists(image_path):
        raise HTTPException(status_code=404, detail="Image not found")

    os.remove(image_path)
    available_images.pop(image_name, None)
    return JSONResponse(content={"message": "Image deleted successfully"})


# Gradio interface function
def upload_image_interface(url: str, file: gr.inputs.File):
    response = upload_image(url, file)
    return response.content

# Gradio interface
gr.Interface(
    fn=upload_image_interface,
    inputs=["text", "file"],
    outputs="text",
    title="Image Uploader",
    description="Upload an image with a URL and file",
    examples=[["https://example.com/image.jpg", gr.inputs.File(file_type="image")]]
).launch()