agent/tools/media_tools.py

"""
NeuraPrompt Agent — Media Tools (v9.0)
Charts, QR codes, image manipulation, screenshots — all with download URLs.
Auto-installs required packages if missing.
"""

import os
import uuid
import subprocess
import logging
from pathlib import Path
from typing import List, Dict, Optional, Any

log = logging.getLogger("agent.tools.media")

OUTPUT_DIR = Path("/tmp/agent_outputs")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

BASE_URL = os.getenv("SPACE_HOST", "https://deepimagix-self-trained2.hf.space")


def _url(filename: str) -> str:
    return f"{BASE_URL}/agent/download/{filename}"


def _install(package: str):
    subprocess.run(["pip", "install", package, "--quiet"], capture_output=True)


# ─────────────────────────────────────────────────────────────
#  BAR CHART
# ─────────────────────────────────────────────────────────────
def create_bar_chart(
    labels: List[str],
    values: List[float],
    title: str = "Chart",
    x_label: str = "",
    y_label: str = "",
    filename: Optional[str] = None
) -> str:
    """
    Create a bar chart image (PNG).

    Args:
        labels   : List of bar labels
        values   : List of numeric values
        title    : Chart title
        x_label  : X axis label
        y_label  : Y axis label
        filename : Optional output filename

    Returns:
        str: Download URL for the chart image
    """
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        _install("matplotlib")
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

    try:
        fname = filename or f"chart_{uuid.uuid4().hex[:8]}.png"
        if not fname.endswith(".png"):
            fname += ".png"
        out_path = OUTPUT_DIR / fname

        fig, ax = plt.subplots(figsize=(10, 6))
        bars = ax.bar(labels, values, color="black", edgecolor="white")
        ax.set_title(title, fontsize=16, fontweight="bold")
        ax.set_xlabel(x_label, fontsize=12)
        ax.set_ylabel(y_label, fontsize=12)
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)

        # Value labels on bars
        for bar, val in zip(bars, values):
            ax.text(
                bar.get_x() + bar.get_width() / 2,
                bar.get_height() + max(values) * 0.01,
                str(round(val, 2)),
                ha="center", va="bottom", fontsize=10
            )

        plt.tight_layout()
        plt.savefig(str(out_path), dpi=150, bbox_inches="tight")
        plt.close()
        return f"✅ Bar chart created!\n📥 Download: {_url(fname)}"
    except Exception as e:
        return f"Bar chart error: {str(e)}"


# ─────────────────────────────────────────────────────────────
#  LINE CHART
# ─────────────────────────────────────────────────────────────
def create_line_chart(
    x_values: List[Any],
    y_values: List[float],
    title: str = "Line Chart",
    x_label: str = "",
    y_label: str = "",
    filename: Optional[str] = None
) -> str:
    """
    Create a line chart image (PNG).

    Args:
        x_values : X axis values (dates, labels, numbers)
        y_values : Y axis numeric values
        title    : Chart title
        x_label  : X axis label
        y_label  : Y axis label
        filename : Optional output filename

    Returns:
        str: Download URL
    """
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        _install("matplotlib")
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

    try:
        fname = filename or f"linechart_{uuid.uuid4().hex[:8]}.png"
        if not fname.endswith(".png"):
            fname += ".png"
        out_path = OUTPUT_DIR / fname

        fig, ax = plt.subplots(figsize=(10, 6))
        ax.plot(x_values, y_values, color="black", linewidth=2, marker="o", markersize=5)
        ax.fill_between(range(len(x_values)), y_values, alpha=0.1, color="black")
        ax.set_title(title, fontsize=16, fontweight="bold")
        ax.set_xlabel(x_label, fontsize=12)
        ax.set_ylabel(y_label, fontsize=12)
        ax.set_xticks(range(len(x_values)))
        ax.set_xticklabels([str(x) for x in x_values], rotation=45, ha="right")
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        plt.tight_layout()
        plt.savefig(str(out_path), dpi=150, bbox_inches="tight")
        plt.close()
        return f"✅ Line chart created!\n📥 Download: {_url(fname)}"
    except Exception as e:
        return f"Line chart error: {str(e)}"


# ─────────────────────────────────────────────────────────────
#  PIE CHART
# ─────────────────────────────────────────────────────────────
def create_pie_chart(
    labels: List[str],
    values: List[float],
    title: str = "Pie Chart",
    filename: Optional[str] = None
) -> str:
    """
    Create a pie chart image (PNG).

    Args:
        labels   : Slice labels
        values   : Slice values
        title    : Chart title
        filename : Optional output filename

    Returns:
        str: Download URL
    """
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        _install("matplotlib")
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

    try:
        fname = filename or f"piechart_{uuid.uuid4().hex[:8]}.png"
        if not fname.endswith(".png"):
            fname += ".png"
        out_path = OUTPUT_DIR / fname

        greys = [str(round(i / len(labels), 2)) for i in range(len(labels))]
        fig, ax = plt.subplots(figsize=(8, 8))
        wedges, texts, autotexts = ax.pie(
            values, labels=labels, autopct="%1.1f%%",
            colors=greys, startangle=90,
            wedgeprops={"edgecolor": "white", "linewidth": 2}
        )
        ax.set_title(title, fontsize=16, fontweight="bold")
        plt.tight_layout()
        plt.savefig(str(out_path), dpi=150, bbox_inches="tight")
        plt.close()
        return f"✅ Pie chart created!\n📥 Download: {_url(fname)}"
    except Exception as e:
        return f"Pie chart error: {str(e)}"


# ─────────────────────────────────────────────────────────────
#  QR CODE
# ─────────────────────────────────────────────────────────────
def create_qr_code(
    data: str,
    filename: Optional[str] = None
) -> str:
    """
    Generate a QR code image (PNG).

    Args:
        data     : Text or URL to encode in the QR code
        filename : Optional output filename

    Returns:
        str: Download URL for the QR code image
    """
    try:
        import qrcode
    except ImportError:
        _install("qrcode[pil]")
        import qrcode

    try:
        fname = filename or f"qrcode_{uuid.uuid4().hex[:8]}.png"
        if not fname.endswith(".png"):
            fname += ".png"
        out_path = OUTPUT_DIR / fname

        qr = qrcode.QRCode(
            version=1,
            error_correction=qrcode.constants.ERROR_CORRECT_H,
            box_size=10,
            border=4,
        )
        qr.add_data(data)
        qr.make(fit=True)
        img = qr.make_image(fill_color="black", back_color="white")
        img.save(str(out_path))
        return f"✅ QR code created for: {data[:50]}\n📥 Download: {_url(fname)}"
    except Exception as e:
        return f"QR code error: {str(e)}"


# ─────────────────────────────────────────────────────────────
#  IMAGE RESIZE / CONVERT
# ─────────────────────────────────────────────────────────────
def resize_image(
    image_b64: str,
    width: int,
    height: int,
    filename: Optional[str] = None
) -> str:
    """
    Resize an image from base64 input.

    Args:
        image_b64 : Base64 encoded image
        width     : Target width in pixels
        height    : Target height in pixels
        filename  : Optional output filename

    Returns:
        str: Download URL for the resized image
    """
    try:
        from PIL import Image
        import io, base64
    except ImportError:
        _install("Pillow")
        from PIL import Image
        import io, base64

    try:
        fname = filename or f"resized_{uuid.uuid4().hex[:8]}.png"
        if not fname.endswith(".png"):
            fname += ".png"
        out_path = OUTPUT_DIR / fname

        img_bytes = base64.b64decode(image_b64)
        img = Image.open(io.BytesIO(img_bytes))
        img = img.resize((width, height), Image.LANCZOS)
        img.save(str(out_path))
        return f"✅ Image resized to {width}x{height}\n📥 Download: {_url(fname)}"
    except Exception as e:
        return f"Image resize error: {str(e)}"


# ─────────────────────────────────────────────────────────────
#  ZIP CREATOR
# ─────────────────────────────────────────────────────────────
def create_zip(
    file_paths: List[str],
    zip_filename: Optional[str] = None
) -> str:
    """
    Create a ZIP archive from a list of files in the output folder.

    Args:
        file_paths   : List of filenames in the agent_outputs folder to zip
        zip_filename : Optional output zip filename

    Returns:
        str: Download URL for the ZIP file
    """
    import zipfile

    try:
        fname = zip_filename or f"archive_{uuid.uuid4().hex[:8]}.zip"
        if not fname.endswith(".zip"):
            fname += ".zip"
        out_path = OUTPUT_DIR / fname

        with zipfile.ZipFile(str(out_path), "w", zipfile.ZIP_DEFLATED) as zf:
            for fp in file_paths:
                full = OUTPUT_DIR / fp
                if full.exists():
                    zf.write(str(full), fp)
                else:
                    log.warning(f"File not found for zip: {fp}")

        return f"✅ ZIP archive created with {len(file_paths)} files\n📥 Download: {_url(fname)}"
    except Exception as e:
        return f"ZIP creation error: {str(e)}"