"""
Armory Media Types - Image, Audio, Video, and Histogram logging

This module provides media logging capabilities similar to W&B.
"""

import os
import io
import base64
import json
from typing import Any, Dict, List, Optional, Union, Sequence
from pathlib import Path
import numpy as np


class Image:
    """
    Image media type for logging images to Armory.

    Similar to wandb.Image, supports numpy arrays, PIL images,
    file paths, and matplotlib figures.

    Example:
        import armory
        from armory import Image

        run = armory.init(project="my-project")

        # Log numpy array
        run.log({"image": Image(np.random.rand(100, 100, 3))})

        # Log PIL image
        run.log({"photo": Image(pil_image, caption="My photo")})

        # Log file path
        run.log({"sample": Image("path/to/image.png")})

        # Log matplotlib figure
        run.log({"plot": Image(plt)})
    """

    def __init__(
        self,
        data_or_path: Any,
        caption: Optional[str] = None,
        mode: Optional[str] = None,
        grouping: Optional[int] = None,
    ):
        """
        Args:
            data_or_path: Image data - numpy array, PIL Image, file path, or matplotlib figure
            caption: Optional caption for the image
            mode: Image mode (e.g., "RGB", "L")
            grouping: Group images together in the UI
        """
        self.caption = caption
        self.mode = mode
        self.grouping = grouping
        self._data = None
        self._path = None
        self._format = "png"

        if isinstance(data_or_path, str):
            self._path = data_or_path
        elif isinstance(data_or_path, Path):
            self._path = str(data_or_path)
        else:
            self._data = data_or_path

    def to_dict(self) -> Dict[str, Any]:
        """Convert image to a serializable dictionary"""
        result = {
            "_type": "image",
            "caption": self.caption,
        }

        if self._path:
            result["path"] = self._path
            # Read and encode file
            if os.path.exists(self._path):
                with open(self._path, "rb") as f:
                    result["data"] = base64.b64encode(f.read()).decode("utf-8")
                    result["format"] = Path(self._path).suffix.lstrip(".")
        elif self._data is not None:
            result["data"] = self._encode_data()
            result["format"] = self._format

        return result

    def _encode_data(self) -> str:
        """Encode image data to base64"""
        data = self._data

        # Handle matplotlib figure
        if hasattr(data, 'savefig'):
            buf = io.BytesIO()
            data.savefig(buf, format='png', bbox_inches='tight')
            buf.seek(0)
            return base64.b64encode(buf.read()).decode("utf-8")

        # Handle PIL Image
        try:
            from PIL import Image as PILImage
            if isinstance(data, PILImage.Image):
                buf = io.BytesIO()
                data.save(buf, format='PNG')
                buf.seek(0)
                return base64.b64encode(buf.read()).decode("utf-8")
        except ImportError:
            pass

        # Handle numpy array
        if hasattr(data, 'shape'):
            try:
                from PIL import Image as PILImage
                if data.ndim == 2:
                    img = PILImage.fromarray(data.astype('uint8'), mode='L')
                elif data.ndim == 3:
                    if data.shape[2] == 4:
                        img = PILImage.fromarray(data.astype('uint8'), mode='RGBA')
                    else:
                        img = PILImage.fromarray(data.astype('uint8'), mode='RGB')
                else:
                    return ""
                buf = io.BytesIO()
                img.save(buf, format='PNG')
                buf.seek(0)
                return base64.b64encode(buf.read()).decode("utf-8")
            except ImportError:
                # Fallback: just store the shape info
                return json.dumps({"shape": list(data.shape), "dtype": str(data.dtype)})

        return ""


class Audio:
    """
    Audio media type for logging audio to Armory.

    Similar to wandb.Audio.

    Example:
        from armory import Audio

        run.log({"audio": Audio(audio_array, sample_rate=16000)})
        run.log({"speech": Audio("path/to/audio.wav", caption="Sample")})
    """

    def __init__(
        self,
        data_or_path: Any,
        sample_rate: int = 16000,
        caption: Optional[str] = None,
    ):
        """
        Args:
            data_or_path: Audio data - numpy array or file path
            sample_rate: Sample rate in Hz
            caption: Optional caption
        """
        self.sample_rate = sample_rate
        self.caption = caption
        self._data = None
        self._path = None

        if isinstance(data_or_path, (str, Path)):
            self._path = str(data_or_path)
        else:
            self._data = data_or_path

    def to_dict(self) -> Dict[str, Any]:
        """Convert audio to a serializable dictionary"""
        result = {
            "_type": "audio",
            "sample_rate": self.sample_rate,
            "caption": self.caption,
        }

        if self._path:
            result["path"] = self._path
            if os.path.exists(self._path):
                with open(self._path, "rb") as f:
                    result["data"] = base64.b64encode(f.read()).decode("utf-8")
                    result["format"] = Path(self._path).suffix.lstrip(".")
        elif self._data is not None:
            # Try to encode as WAV
            try:
                import scipy.io.wavfile as wav
                buf = io.BytesIO()
                wav.write(buf, self.sample_rate, self._data)
                buf.seek(0)
                result["data"] = base64.b64encode(buf.read()).decode("utf-8")
                result["format"] = "wav"
            except ImportError:
                result["data"] = base64.b64encode(self._data.tobytes()).decode("utf-8")
                result["format"] = "raw"

        return result


class Video:
    """
    Video media type for logging videos to Armory.

    Similar to wandb.Video.

    Example:
        from armory import Video

        run.log({"video": Video(frames_array, fps=30)})
        run.log({"recording": Video("path/to/video.mp4")})
    """

    def __init__(
        self,
        data_or_path: Any,
        fps: int = 30,
        format: str = "mp4",
        caption: Optional[str] = None,
    ):
        """
        Args:
            data_or_path: Video data - numpy array (T, H, W, C) or file path
            fps: Frames per second
            format: Video format
            caption: Optional caption
        """
        self.fps = fps
        self.format = format
        self.caption = caption
        self._data = None
        self._path = None

        if isinstance(data_or_path, (str, Path)):
            self._path = str(data_or_path)
        else:
            self._data = data_or_path

    def to_dict(self) -> Dict[str, Any]:
        """Convert video to a serializable dictionary"""
        result = {
            "_type": "video",
            "fps": self.fps,
            "format": self.format,
            "caption": self.caption,
        }

        if self._path:
            result["path"] = self._path
            if os.path.exists(self._path):
                with open(self._path, "rb") as f:
                    result["data"] = base64.b64encode(f.read()).decode("utf-8")
        elif self._data is not None:
            # Store frame info (actual video encoding would require moviepy/ffmpeg)
            result["frames"] = len(self._data) if hasattr(self._data, '__len__') else 0
            if hasattr(self._data, 'shape'):
                result["shape"] = list(self._data.shape)

        return result


class Histogram:
    """
    Histogram data type for logging distributions to Armory.

    Similar to wandb.Histogram.

    Example:
        from armory import Histogram

        run.log({"activations": Histogram(layer_outputs)})
        run.log({"weights": Histogram(model.weights, num_bins=50)})
    """

    def __init__(
        self,
        sequence: Sequence[float],
        num_bins: int = 64,
    ):
        """
        Args:
            sequence: Sequence of values to histogram
            num_bins: Number of histogram bins
        """
        self.num_bins = num_bins
        self._sequence = sequence

    def to_dict(self) -> Dict[str, Any]:
        """Convert histogram to a serializable dictionary"""
        try:
            data = np.array(self._sequence).flatten()
            counts, bin_edges = np.histogram(data, bins=self.num_bins)
            return {
                "_type": "histogram",
                "counts": counts.tolist(),
                "bin_edges": bin_edges.tolist(),
                "min": float(data.min()),
                "max": float(data.max()),
                "mean": float(data.mean()),
                "std": float(data.std()),
                "count": len(data),
            }
        except Exception:
            return {
                "_type": "histogram",
                "error": "Failed to compute histogram",
            }


class Table:
    """
    Table data type for logging tabular data to Armory.

    Similar to wandb.Table.

    Example:
        from armory import Table

        table = Table(columns=["id", "image", "label"])
        table.add_data(1, Image(img1), "cat")
        table.add_data(2, Image(img2), "dog")
        run.log({"predictions": table})
    """

    def __init__(
        self,
        columns: Optional[List[str]] = None,
        data: Optional[List[List[Any]]] = None,
        dataframe: Any = None,
    ):
        """
        Args:
            columns: List of column names
            data: List of rows (each row is a list of values)
            dataframe: Pandas DataFrame to convert
        """
        self.columns = columns or []
        self.data = data or []

        if dataframe is not None:
            self._from_dataframe(dataframe)

    def _from_dataframe(self, df):
        """Initialize from a pandas DataFrame"""
        self.columns = list(df.columns)
        self.data = df.values.tolist()

    def add_data(self, *row):
        """Add a row of data"""
        self.data.append(list(row))

    def add_column(self, name: str, data: List[Any]):
        """Add a column"""
        self.columns.append(name)
        for i, val in enumerate(data):
            if i < len(self.data):
                self.data[i].append(val)
            else:
                self.data.append([None] * (len(self.columns) - 1) + [val])

    def to_dict(self) -> Dict[str, Any]:
        """Convert table to a serializable dictionary"""
        # Convert any media objects in data
        serialized_data = []
        for row in self.data:
            serialized_row = []
            for val in row:
                if hasattr(val, 'to_dict'):
                    serialized_row.append(val.to_dict())
                else:
                    serialized_row.append(val)
            serialized_data.append(serialized_row)

        return {
            "_type": "table",
            "columns": self.columns,
            "data": serialized_data,
            "row_count": len(self.data),
        }


class Html:
    """
    HTML content type for logging HTML to Armory.

    Example:
        from armory import Html

        run.log({"report": Html("<h1>Results</h1><p>Accuracy: 95%</p>")})
    """

    def __init__(self, data: str, inject: bool = True):
        """
        Args:
            data: HTML string or file path
            inject: Whether to inject into the dashboard
        """
        self.inject = inject

        if os.path.isfile(data):
            with open(data, "r") as f:
                self._html = f.read()
        else:
            self._html = data

    def to_dict(self) -> Dict[str, Any]:
        return {
            "_type": "html",
            "html": self._html,
            "inject": self.inject,
        }


class Object3D:
    """
    3D object data type for logging 3D models to Armory.

    Example:
        from armory import Object3D

        run.log({"model_3d": Object3D("path/to/model.obj")})
    """

    def __init__(
        self,
        data_or_path: Any,
        caption: Optional[str] = None,
    ):
        self.caption = caption
        self._path = None
        self._data = None

        if isinstance(data_or_path, (str, Path)):
            self._path = str(data_or_path)
        else:
            self._data = data_or_path

    def to_dict(self) -> Dict[str, Any]:
        result = {
            "_type": "object3d",
            "caption": self.caption,
        }

        if self._path:
            result["path"] = self._path
            if os.path.exists(self._path):
                with open(self._path, "rb") as f:
                    result["data"] = base64.b64encode(f.read()).decode("utf-8")
                    result["format"] = Path(self._path).suffix.lstrip(".")

        return result
