agent/tools/document_tools.py

"""
NeuraPrompt Agent — Document Tools (v9.0)
Generates PDF, DOCX, Excel, CSV, TXT files 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.document")

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)


# ─────────────────────────────────────────────────────────────
#  PDF
# ─────────────────────────────────────────────────────────────
def create_pdf(
    title: str,
    content: str,
    filename: Optional[str] = None
) -> str:
    """
    Create a PDF document.

    Args:
        title    : Document title shown at top
        content  : Body text (use \\n for new lines)
        filename : Optional output filename (auto-generated if not set)

    Returns:
        str: Download URL for the PDF
    """
    try:
        from reportlab.lib.pagesizes import A4
        from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
        from reportlab.lib.units import cm
        from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
        from reportlab.lib.enums import TA_LEFT
    except ImportError:
        _install("reportlab")
        from reportlab.lib.pagesizes import A4
        from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
        from reportlab.lib.units import cm
        from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
        from reportlab.lib.enums import TA_LEFT

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

        doc = SimpleDocTemplate(
            str(out_path),
            pagesize=A4,
            leftMargin=2*cm, rightMargin=2*cm,
            topMargin=2*cm, bottomMargin=2*cm
        )
        styles = getSampleStyleSheet()
        story = []

        # Title
        title_style = ParagraphStyle(
            "Title", parent=styles["Heading1"],
            fontSize=18, spaceAfter=12
        )
        story.append(Paragraph(title, title_style))
        story.append(Spacer(1, 0.5*cm))

        # Body
        body_style = ParagraphStyle(
            "Body", parent=styles["Normal"],
            fontSize=11, leading=16
        )
        for line in content.split("\n"):
            if line.strip():
                story.append(Paragraph(line.strip(), body_style))
                story.append(Spacer(1, 0.2*cm))

        doc.build(story)
        return f"✅ PDF created!\n📥 Download: {_url(fname)}"
    except Exception as e:
        return f"PDF creation error: {str(e)}"


# ─────────────────────────────────────────────────────────────
#  DOCX
# ─────────────────────────────────────────────────────────────
def create_docx(
    title: str,
    content: str,
    filename: Optional[str] = None,
    sections: Optional[List[Dict[str, str]]] = None
) -> str:
    """
    Create a Word document (.docx).

    Args:
        title    : Document title
        content  : Main body text (use \\n for new lines)
        filename : Optional output filename
        sections : Optional list of {heading, text} dicts for structured docs

    Returns:
        str: Download URL for the DOCX
    """
    try:
        from docx import Document
        from docx.shared import Pt, Inches
    except ImportError:
        _install("python-docx")
        from docx import Document
        from docx.shared import Pt, Inches

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

        doc = Document()

        # Title
        title_para = doc.add_heading(title, level=0)
        title_para.runs[0].font.size = Pt(20)

        # Sections if provided
        if sections:
            for section in sections:
                heading = section.get("heading", "")
                text    = section.get("text", "")
                if heading:
                    doc.add_heading(heading, level=1)
                if text:
                    for line in text.split("\n"):
                        if line.strip():
                            doc.add_paragraph(line.strip())
        else:
            # Plain content
            for line in content.split("\n"):
                if line.strip():
                    doc.add_paragraph(line.strip())

        doc.save(str(out_path))
        return f"✅ DOCX created!\n📥 Download: {_url(fname)}"
    except Exception as e:
        return f"DOCX creation error: {str(e)}"


# ─────────────────────────────────────────────────────────────
#  EXCEL
# ─────────────────────────────────────────────────────────────
def create_excel(
    sheet_name: str,
    headers: List[str],
    rows: List[List[Any]],
    filename: Optional[str] = None
) -> str:
    """
    Create an Excel spreadsheet (.xlsx).

    Args:
        sheet_name : Name of the worksheet
        headers    : List of column header names
        rows       : List of rows, each row is a list of values
        filename   : Optional output filename

    Returns:
        str: Download URL for the Excel file

    Example:
        create_excel(
            sheet_name="Sales",
            headers=["Name", "Amount", "Date"],
            rows=[["Alice", 1000, "2026-06-01"], ["Bob", 2000, "2026-06-02"]]
        )
    """
    try:
        import openpyxl
        from openpyxl.styles import Font, PatternFill, Alignment
    except ImportError:
        _install("openpyxl")
        import openpyxl
        from openpyxl.styles import Font, PatternFill, Alignment

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

        wb = openpyxl.Workbook()
        ws = wb.active
        ws.title = sheet_name

        # Header row with bold styling
        header_fill = PatternFill("solid", fgColor="000000")
        header_font = Font(bold=True, color="FFFFFF")
        for col, header in enumerate(headers, 1):
            cell = ws.cell(row=1, column=col, value=header)
            cell.font = header_font
            cell.fill = header_fill
            cell.alignment = Alignment(horizontal="center")
            ws.column_dimensions[cell.column_letter].width = max(15, len(str(header)) + 4)

        # Data rows
        for r_idx, row in enumerate(rows, 2):
            for c_idx, value in enumerate(row, 1):
                ws.cell(row=r_idx, column=c_idx, value=value)

        wb.save(str(out_path))
        return f"✅ Excel file created! {len(rows)} rows × {len(headers)} columns\n📥 Download: {_url(fname)}"
    except Exception as e:
        return f"Excel creation error: {str(e)}"


# ─────────────────────────────────────────────────────────────
#  CSV
# ─────────────────────────────────────────────────────────────
def create_csv(
    headers: List[str],
    rows: List[List[Any]],
    filename: Optional[str] = None
) -> str:
    """
    Create a CSV file.

    Args:
        headers  : Column header names
        rows     : List of rows
        filename : Optional output filename

    Returns:
        str: Download URL for the CSV
    """
    import csv

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

        with open(out_path, "w", newline="", encoding="utf-8") as f:
            writer = csv.writer(f)
            writer.writerow(headers)
            writer.writerows(rows)

        return f"✅ CSV created! {len(rows)} rows\n📥 Download: {_url(fname)}"
    except Exception as e:
        return f"CSV creation error: {str(e)}"


# ─────────────────────────────────────────────────────────────
#  TXT / MARKDOWN
# ─────────────────────────────────────────────────────────────
def create_text_file(
    content: str,
    filename: Optional[str] = None,
    file_format: str = "txt"
) -> str:
    """
    Create a plain text or markdown file.

    Args:
        content     : File content
        filename    : Optional output filename
        file_format : "txt" or "md"

    Returns:
        str: Download URL
    """
    try:
        ext   = file_format.lstrip(".")
        fname = filename or f"file_{uuid.uuid4().hex[:8]}.{ext}"
        if not fname.endswith(f".{ext}"):
            fname += f".{ext}"
        out_path = OUTPUT_DIR / fname
        out_path.write_text(content, encoding="utf-8")
        return f"✅ {ext.upper()} file created!\n📥 Download: {_url(fname)}"
    except Exception as e:
        return f"Text file creation error: {str(e)}"