from base64 import b64decode, b64encode
from io import BytesIO
import torch
from PIL import Image
import numpy as np
import torchvision.transforms as transforms
import subprocess
import re
from imageio_ffmpeg import get_ffmpeg_exe

FFMPEG_PATH = get_ffmpeg_exe()
ENCODE_ARGS = ("utf-8", 'backslashreplace')

def convert_base64_to_image(base64_str: str) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Decodes a Base64-encoded image into a ComfyUI image tensor and an alpha mask.
    """
    try:
        image = Image.open(BytesIO(b64decode(base64_str)))
        W, H = image.size

        if "A" in image.getbands():
            alpha = np.array(image.getchannel("A")).astype(np.float32) / 255.0
            mask = torch.from_numpy(1.0 - alpha)
        else:
            mask = torch.zeros((H, W), dtype=torch.float32)

        image = image.convert("RGB")
        image = np.array(image).astype(np.float32) / 255.0
        image = torch.from_numpy(image).contiguous()[None,]

        return image, mask
    except Exception as e:
        raise RuntimeError(f"Failed to decode image from Base64 string: {e}")
    
def convert_image_to_base64(image_tensor: torch.Tensor) -> str:
    """
    Encodes a PyTorch image tensor into a Base64-encoded image.
    """
    try:
        # Handle batch dimension
        if image_tensor.ndim == 4:
            image_tensor = image_tensor[0]

        print(image_tensor.shape)

        # Ensure (C, H, W) format
        if image_tensor.ndim == 3 and image_tensor.shape[2] in [1, 3, 4]:
            image_tensor = image_tensor.permute(2, 0, 1)

        print(image_tensor.shape)

        # Normalize pixel values (0.0, 1.0)
        image_tensor = torch.clamp(image_tensor, 0.0, 1.0)

        image = transforms.ToPILImage()(image_tensor)

        buffer = BytesIO()
        image.save(buffer, format="PNG")

        return b64encode(buffer.getvalue()).decode("utf-8")
    except Exception as e:
        raise RuntimeError(f"Failed to encode image into Base64 string: {e}")

def convert_base64_to_audio(base64_str: str, start_time=0, duration=0):
    """
    Decodes a Base64-encoded audio file into a PyTorch waveform using FFmpeg.
    """
    try:
        args = [FFMPEG_PATH, '-i', 'pipe:0', '-vn', '-acodec', 'pcm_f32le', '-f', 'f32le', 'pipe:1']

        if start_time > 0:
            args.insert(1, '-ss')
            args.insert(2, str(start_time))
        if duration > 0:
            args.insert(1, '-t')
            args.insert(2, str(duration))

        try:
            audio_data = b64decode(base64_str)
            res = subprocess.run(args, input=audio_data, capture_output=True, check=True)

            audio = torch.frombuffer(res.stdout, dtype=torch.float32)

            match = re.search(r', (\d+) Hz, (\w+),', res.stderr.decode(*ENCODE_ARGS))
            if match:
                sample_rate = int(match.group(1))
                channels = {"mono": 1, "stereo": 2}[match.group(2)]
            else:
                sample_rate = 44100
                channels = 2

            audio = audio.reshape((-1, channels)).transpose(0, 1).unsqueeze(0)

            return {'waveform': audio, 'sample_rate': sample_rate}
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"FFmpeg failed:\n{e.stderr.decode(*ENCODE_ARGS)}")
    except Exception as e:
        raise RuntimeError(f"Failed to decode audio: {e}")