import base64
from io import BytesIO
from PIL import Image

def decode_chart(charts: dict):
    """
    Decodes base64 string charts into PIL Image objects.
    Also converts Matplotlib Figures into PIL Image objects.
    If already decoded or None, returns as is.
    """
    if not isinstance(charts, dict):
        return {}
        
    result = {}
    for name, img_str in charts.items():
        # Case 1: Base64 String (from DB)
        if isinstance(img_str, str) and not img_str.startswith("<matplotlib"):
            try:
                # Remove data URI prefix if present
                if "," in img_str:
                    img_str = img_str.split(",")[1]
                img_bytes = base64.b64decode(img_str)
                img = Image.open(BytesIO(img_bytes))
                result[name] = img
            except Exception as e:
                import logging
                logging.getLogger(__name__).error(f"Failed to decode base64 chart {name}: {e}")
                result[name] = img_str
        # Case 2: Matplotlib Figure (from fresh run)
        elif hasattr(img_str, "savefig"):
            try:
                buf = BytesIO()
                img_str.savefig(buf, format="png", bbox_inches='tight')
                buf.seek(0)
                img = Image.open(buf)
                result[name] = img
            except Exception as e:
                import logging
                logging.getLogger(__name__).error(f"Failed to convert figure to image for {name}: {e}")
                result[name] = img_str
        else:
            result[name] = img_str

    return result

def figure_to_base64(fig):
    """
    Converts a Matplotlib Figure to a base64 string for database storage.
    """
    buf = BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight')
    buf.seek(0)
    img_str = base64.b64encode(buf.getvalue()).decode('utf-8')
    return img_str

def serialize_charts(charts: dict) -> dict:
    """
    Ensures all charts in the dict are JSON-serializable base64 strings.
    Handles:
    - str: Assumed to be already base64.
    - plt.Figure: Serialized using savefig.
    - PIL.Image: Serialized using img.save.
    """
    if not isinstance(charts, dict):
        return {}
    
    result = {}
    for name, val in charts.items():
        if val is None:
            continue
            
        if isinstance(val, str):
            # already serialized
            result[name] = val
        elif hasattr(val, "savefig"):
            # Matplotlib Figure
            result[name] = figure_to_base64(val)
        elif hasattr(val, "save"):
            # PIL Image
            try:
                buf = BytesIO()
                val.save(buf, format="png")
                buf.seek(0)
                result[name] = base64.b64encode(buf.getvalue()).decode("utf-8")
            except Exception as e:
                import logging
                logging.getLogger(__name__).error(f"Failed to serialize PIL image for {name}: {e}")
                result[name] = None
        else:
            result[name] = str(val)

    return result