from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import StreamingResponse
from image_enhancer import EnhancementMethod, Enhancer
from video_enhancer import VideoEnhancer
from pydantic import BaseModel
from PIL import Image
from io import BytesIO
import base64
import magic
from typing import List
class EnhancementRequest(BaseModel):
    method: EnhancementMethod = EnhancementMethod.gfpgan
    background_enhancement: bool = True
    upscale: int = 2

class _EnhanceBase(BaseModel):
    encoded_base_img: List[str]



app = FastAPI()

@app.get("/")
def greet_json():
    return {"Initializing GlamApp Enhancer"}

@app.post("/enhance/image/")
async def enhance_image(
    file: UploadFile = File(...),
    method: EnhancementMethod = EnhancementMethod.gfpgan,
    background_enhancement: bool = True,
    upscale: int = 2
):
    try:
        if not file.content_type.startswith('image/'):
            raise HTTPException(status_code=400, detail="Invalid file type")
        
        contents = await file.read()
        base64_encoded_image = base64.b64encode(contents).decode('utf-8')
        
        request = EnhancementRequest(
            method=method,
            background_enhancement=background_enhancement,
            upscale=upscale
        )
        
        enhancer = Enhancer(request.method, request.background_enhancement, request.upscale)
        enhanced_img, original_resolution, enhanced_resolution = await enhancer.enhance(base64_encoded_image)
        
        enhanced_image = Image.fromarray(enhanced_img)
        img_byte_arr = BytesIO()
        enhanced_image.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0)
        
        print(f"Original resolution: {original_resolution}, Enhanced resolution: {enhanced_resolution}")
        
        return StreamingResponse(img_byte_arr, media_type="image/png")
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
        
@app.post("/enhance/video/")
async def enhance_video(file: UploadFile = File(...)):
    enhancer = VideoEnhancer()
    file_header = await file.read(1024)
    file.file.seek(0)
    mime = magic.Magic(mime=True)
    file_mime_type = mime.from_buffer(file_header)
    
    accepted_mime_types = [
        'video/mp4',
        'video/mpeg',
        'video/x-msvideo',
        'video/quicktime',
        'video/x-matroska',
        'video/webm'
    ]
    
    if file_mime_type not in accepted_mime_types:
        raise HTTPException(status_code=400, detail="Invalid file type. Please upload a video file.")
    return await enhancer.stream_enhanced_video(file)