from PIL import Image, ImageDraw, ImageFont
import os
from typing import Dict, Any
import cv2
import numpy as np
import binascii
from typing import Union
import base64
import io

def add_watermark(image_input: Union[str, Image.Image], watermark_text: str, opacity: float = 0.5) -> Union[str, Image.Image]:
    """
    Add a semi-transparent text watermark to an image.
    
    Args:
        image_input: PIL Image object or base64 string
        watermark_text: The text to be used as watermark
        opacity: Opacity of the watermark (0.1 to 1.0)
        
    Returns:
        PIL Image object with watermark applied
    """
    try:
        if isinstance(image_input, str):
            if image_input.startswith('data:image'):
                base64_data = image_input.split(',')[1]
            else:
                base64_data = image_input
            image_data = base64.b64decode(base64_data)
            image = Image.open(io.BytesIO(image_data))
        else:
            image = image_input
        
        if image.mode != 'RGBA':
            image = image.convert('RGBA')
        
        overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
        draw = ImageDraw.Draw(overlay)
        
        try:
            font_size = min(image.width, image.height) // 20
            font = ImageFont.truetype("arial.ttf", font_size)
        except:
            try:
                font = ImageFont.load_default()
            except:
                font = ImageFont.load_default()
        
        bbox = draw.textbbox((0, 0), watermark_text, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]
        
        x = (image.width - text_width) // 2
        y = (image.height - text_height) // 2
        
        alpha_value = int(255 * opacity)
        text_color = (255, 255, 255, alpha_value)
        shadow_color = (0, 0, 0, int(alpha_value * 0.5))
        
        draw.text((x-2, y-2), watermark_text, fill=shadow_color, font=font)
        draw.text((x, y), watermark_text, fill=text_color, font=font)
        
        watermarked = Image.alpha_composite(image, overlay)
        final_image = watermarked.convert('RGB')
        
        return final_image
        
    except Exception as e:
        print(f"Error adding watermark: {e}")
        return image_input if isinstance(image_input, Image.Image) else None

def remove_watermark(image_input: Union[str, Image.Image], alpha: float = 2.0, beta: float = -160) -> Union[str, Image.Image]:
    try:
        if isinstance(image_input, str):
            if image_input.startswith('data:image'):
                base64_data = image_input.split(',')[1]
            else:
                base64_data = image_input
            image_data = base64.b64decode(base64_data)
            image = Image.open(io.BytesIO(image_data))
        else:
            image = image_input
        
        img_array = np.array(image)
        
        new = alpha * img_array + beta
        new = np.clip(new, 0, 255).astype(np.uint8)
        
        result_image = Image.fromarray(new)
        
        return result_image
        
    except Exception as e:
        print(f"Error removing watermark: {e}")
        return image_input if isinstance(image_input, Image.Image) else None