from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
import cv2
import numpy as np
from PIL import Image
import noise
import io
import base64
from pydantic import BaseModel

app = FastAPI(title="Advanced Material Map Generator API")

# Request model for input
class MapRequest(BaseModel):
    image_base64: str
    normal_strength: float = 1.0
    normal_blur: int = 5
    normal_bilateral: bool = False
    normal_color: float = 0.3
    disp_contrast: float = 1.0
    disp_noise: bool = False
    disp_noise_scale: float = 0.1
    disp_edge: float = 1.0
    rough_invert: bool = True
    rough_sharpness: float = 1.0
    rough_detail: float = 0.5
    rough_freq: float = 0.5

def generate_normal_map(image: np.ndarray, strength: float, blur_size: int, use_bilateral: bool, color_influence: float) -> Image.Image:
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    if use_bilateral:
        gray = cv2.bilateralFilter(gray, 9, 75, 75)
    else:
        gray = cv2.GaussianBlur(gray, (blur_size, blur_size), 0)
    
    levels = 3
    normal_map = np.zeros((gray.shape[0], gray.shape[1], 3), dtype=np.float32)
    for i in range(levels):
        scale = 1 / (2 ** i)
        resized = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
        sobel_x = cv2.Scharr(resized, cv2.CV_64F, 1, 0)
        sobel_y = cv2.Scharr(resized, cv2.CV_64F, 0, 1)
        sobel_x = cv2.resize(sobel_x, (gray.shape[1], gray.shape[0]), interpolation=cv2.INTER_LINEAR)
        sobel_y = cv2.resize(sobel_y, (gray.shape[1], gray.shape[0]), interpolation=cv2.INTER_LINEAR)
        normal_map[..., 0] += sobel_x * (1.0 / levels)
        normal_map[..., 1] += sobel_y * (1.0 / levels)
    
    normal_map[..., 0] = cv2.normalize(normal_map[..., 0], None, -strength, strength, cv2.NORM_MINMAX)
    normal_map[..., 1] = cv2.normalize(normal_map[..., 1], None, -strength, strength, cv2.NORM_MINMAX)
    normal_map[..., 2] = 1.0
    
    color_factor = color_influence * strength
    normal_map[..., 0] += (image[..., 0] / 255.0 - 0.5) * color_factor
    normal_map[..., 1] += (image[..., 1] / 255.0 - 0.5) * color_factor
    
    norm = np.linalg.norm(normal_map, axis=2, keepdims=True)
    normal_map = np.divide(normal_map, norm, out=np.zeros_like(normal_map), where=norm != 0)
    normal_map = (normal_map + 1) * 127.5
    normal_map = np.clip(normal_map, 0, 255).astype(np.uint8)
    return Image.fromarray(normal_map)

def generate_displacement_map(image: np.ndarray, contrast: float, add_noise: bool, noise_scale: float, edge_boost: float) -> Image.Image:
    img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    img = clahe.apply(img)
    img = cv2.convertScaleAbs(img, alpha=contrast, beta=0)
    laplacian = cv2.Laplacian(img, cv2.CV_64F)
    laplacian = cv2.convertScaleAbs(laplacian, alpha=edge_boost, beta=0)
    img = cv2.addWeighted(img, 1.0, laplacian, 0.5 * edge_boost, 0)
    if add_noise:
        height, width = img.shape
        noise_map = np.zeros((height, width), dtype=np.float32)
        for y in range(height):
            for x in range(width):
                noise_map[y, x] = noise.pnoise2(x / 50.0, y / 50.0, octaves=6) * noise_scale * 255
        img = cv2.add(img, noise_map.astype(np.uint8))
    return Image.fromarray(img)

def generate_roughness_map(image: np.ndarray, invert: bool, sharpness: float, detail_boost: float, frequency_weight: float) -> Image.Image:
    img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    low_freq = cv2.bilateralFilter(img, 9, 75, 75)
    high_freq = cv2.subtract(img, low_freq)
    img = cv2.addWeighted(low_freq, 1.0 - frequency_weight, high_freq, frequency_weight, 0)
    if invert:
        img = 255 - img
    blurred = cv2.GaussianBlur(img, (5, 5), 0)
    img = cv2.addWeighted(img, 1.0 + sharpness, blurred, -sharpness, 0)
    img = cv2.addWeighted(img, 1.0 + detail_boost, blurred, -detail_boost, 0)
    return Image.fromarray(img)

def image_to_base64(img: Image.Image) -> str:
    buffered = io.BytesIO()
    img.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

@app.post("/generate_maps/")
async def generate_maps(request: MapRequest):
    try:
        # Decode base64 image
        image_bytes = base64.b64decode(request.image_base64)
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        img_array = np.array(image)

        # Generate maps
        normal_map = generate_normal_map(
            img_array, request.normal_strength, request.normal_blur, 
            request.normal_bilateral, request.normal_color
        )
        displacement_map = generate_displacement_map(
            img_array, request.disp_contrast, request.disp_noise, 
            request.disp_noise_scale, request.disp_edge
        )
        roughness_map = generate_roughness_map(
            img_array, request.rough_invert, request.rough_sharpness, 
            request.rough_detail, request.rough_freq
        )

        # Convert to base64
        normal_base64 = image_to_base64(normal_map)
        displacement_base64 = image_to_base64(displacement_map)
        roughness_base64 = image_to_base64(roughness_map)

        return JSONResponse(content={
            "status": "success",
            "normal_map": normal_base64,
            "displacement_map": displacement_base64,
            "roughness_map": roughness_base64
        })

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))