server/document_generator.py

"""
S.T.E.W Document Generator — REAL binary file generation.
Supports: PDF, DOCX, XLSX, PPTX, HTML
All files returned as base64-encoded strings.
"""
import base64
import io
import json
import logging
from datetime import datetime
from typing import Any, Optional

from fastapi import HTTPException

logger = logging.getLogger(__name__)


def _to_base64(buf: io.BytesIO) -> str:
    buf.seek(0)
    return base64.b64encode(buf.read()).decode("utf-8")


# ── PDF ───────────────────────────────────────────────────────────────────────

def generate_pdf(content: str, title: str = "Document") -> dict:
    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, Table, TableStyle
        from reportlab.lib import colors
        from reportlab.lib.enums import TA_CENTER, TA_LEFT

        buf = io.BytesIO()
        doc = SimpleDocTemplate(
            buf,
            pagesize=A4,
            rightMargin=2 * cm,
            leftMargin=2 * cm,
            topMargin=2 * cm,
            bottomMargin=2 * cm,
            title=title,
            author="S.T.E.W Agent",
        )

        styles = getSampleStyleSheet()
        title_style = ParagraphStyle(
            "CustomTitle",
            parent=styles["Title"],
            fontSize=18,
            spaceAfter=12,
            alignment=TA_CENTER,
        )
        body_style = ParagraphStyle(
            "CustomBody",
            parent=styles["Normal"],
            fontSize=11,
            leading=16,
            spaceAfter=8,
        )

        story = []
        story.append(Paragraph(title, title_style))
        story.append(Spacer(1, 0.5 * cm))

        # Parse content — split by newline, detect headings
        for line in content.split("\n"):
            stripped = line.strip()
            if not stripped:
                story.append(Spacer(1, 0.3 * cm))
            elif stripped.startswith("## "):
                story.append(Paragraph(stripped[3:], styles["Heading2"]))
            elif stripped.startswith("# "):
                story.append(Paragraph(stripped[2:], styles["Heading1"]))
            elif stripped.startswith("- ") or stripped.startswith("* "):
                story.append(Paragraph(f"• {stripped[2:]}", body_style))
            else:
                # Escape HTML-like chars for reportlab
                safe = stripped.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
                story.append(Paragraph(safe, body_style))

        # Footer with timestamp
        story.append(Spacer(1, 1 * cm))
        footer_style = ParagraphStyle("Footer", parent=styles["Normal"], fontSize=8,
                                      textColor=colors.grey, alignment=TA_CENTER)
        story.append(Paragraph(
            f"Generated by S.T.E.W Agent • {datetime.utcnow().strftime('%Y-%m-%d %H:%M UTC')}",
            footer_style,
        ))

        doc.build(story)
        filename = f"{title.replace(' ', '_')}_{datetime.utcnow().strftime('%Y%m%d%H%M%S')}.pdf"
        return {
            "file": _to_base64(buf),
            "filename": filename,
            "mime_type": "application/pdf",
            "success": True,
        }
    except ImportError:
        raise HTTPException(500, "reportlab not installed — add it to requirements.txt")
    except Exception as e:
        logger.error(f"PDF generation error: {e}")
        raise HTTPException(500, f"PDF generation failed: {e}")


# ── DOCX ──────────────────────────────────────────────────────────────────────

def generate_docx(content: str, title: str = "Document") -> dict:
    try:
        from docx import Document
        from docx.shared import Pt, RGBColor
        from docx.enum.text import WD_ALIGN_PARAGRAPH

        doc = Document()

        # Document title
        title_para = doc.add_heading(title, level=0)
        title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER

        doc.add_paragraph()  # spacer

        for line in content.split("\n"):
            stripped = line.strip()
            if not stripped:
                doc.add_paragraph()
            elif stripped.startswith("## "):
                doc.add_heading(stripped[3:], level=2)
            elif stripped.startswith("# "):
                doc.add_heading(stripped[2:], level=1)
            elif stripped.startswith("### "):
                doc.add_heading(stripped[4:], level=3)
            elif stripped.startswith("- ") or stripped.startswith("* "):
                doc.add_paragraph(stripped[2:], style="List Bullet")
            elif stripped.startswith("1. ") or (len(stripped) > 2 and stripped[0].isdigit() and stripped[1] == "."):
                doc.add_paragraph(stripped, style="List Number")
            else:
                doc.add_paragraph(stripped)

        # Footer paragraph
        footer = doc.add_paragraph()
        footer.alignment = WD_ALIGN_PARAGRAPH.CENTER
        run = footer.add_run(
            f"Generated by S.T.E.W Agent • {datetime.utcnow().strftime('%Y-%m-%d %H:%M UTC')}"
        )
        run.font.size = Pt(8)
        run.font.color.rgb = RGBColor(0x99, 0x99, 0x99)

        buf = io.BytesIO()
        doc.save(buf)
        filename = f"{title.replace(' ', '_')}_{datetime.utcnow().strftime('%Y%m%d%H%M%S')}.docx"
        return {
            "file": _to_base64(buf),
            "filename": filename,
            "mime_type": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
            "success": True,
        }
    except ImportError:
        raise HTTPException(500, "python-docx not installed")
    except Exception as e:
        logger.error(f"DOCX generation error: {e}")
        raise HTTPException(500, f"DOCX generation failed: {e}")


# ── XLSX ──────────────────────────────────────────────────────────────────────

def generate_xlsx(
    data: list[dict],
    sheet_name: str = "Sheet1",
    title: str = "Spreadsheet",
) -> dict:
    try:
        import openpyxl
        from openpyxl.styles import Font, PatternFill, Alignment, Border, Side
        from openpyxl.utils import get_column_letter

        wb = openpyxl.Workbook()
        ws = wb.active
        ws.title = sheet_name[:31]  # Excel sheet name limit

        if not data:
            data = [{"No data": "provided"}]

        headers = list(data[0].keys())

        # Header row styling
        header_fill = PatternFill(start_color="2563EB", end_color="2563EB", fill_type="solid")
        header_font = Font(color="FFFFFF", bold=True, size=11)
        border = Border(
            left=Side(style="thin"),
            right=Side(style="thin"),
            top=Side(style="thin"),
            bottom=Side(style="thin"),
        )

        for col_idx, header in enumerate(headers, 1):
            cell = ws.cell(row=1, column=col_idx, value=header)
            cell.font = header_font
            cell.fill = header_fill
            cell.alignment = Alignment(horizontal="center", vertical="center")
            cell.border = border

        # Data rows
        alt_fill = PatternFill(start_color="EFF6FF", end_color="EFF6FF", fill_type="solid")
        for row_idx, row in enumerate(data, 2):
            for col_idx, header in enumerate(headers, 1):
                cell = ws.cell(row=row_idx, column=col_idx, value=row.get(header, ""))
                cell.border = border
                if row_idx % 2 == 0:
                    cell.fill = alt_fill

        # Auto-fit columns
        for col_idx, header in enumerate(headers, 1):
            max_len = max(
                len(str(header)),
                *[len(str(row.get(header, ""))) for row in data],
            )
            ws.column_dimensions[get_column_letter(col_idx)].width = min(max_len + 4, 50)

        ws.freeze_panes = "A2"

        buf = io.BytesIO()
        wb.save(buf)
        filename = f"{title.replace(' ', '_')}_{datetime.utcnow().strftime('%Y%m%d%H%M%S')}.xlsx"
        return {
            "file": _to_base64(buf),
            "filename": filename,
            "mime_type": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            "success": True,
        }
    except ImportError:
        raise HTTPException(500, "openpyxl not installed")
    except Exception as e:
        logger.error(f"XLSX generation error: {e}")
        raise HTTPException(500, f"XLSX generation failed: {e}")


# ── PPTX ──────────────────────────────────────────────────────────────────────

def generate_pptx(slides: list[dict], title: str = "Presentation") -> dict:
    try:
        from pptx import Presentation
        from pptx.util import Inches, Pt, Emu
        from pptx.dml.color import RGBColor
        from pptx.enum.text import PP_ALIGN

        prs = Presentation()
        prs.slide_width = Inches(13.33)
        prs.slide_height = Inches(7.5)

        DARK_BG = RGBColor(0x1E, 0x29, 0x3B)
        ACCENT = RGBColor(0x38, 0xBD, 0xF8)
        WHITE = RGBColor(0xFF, 0xFF, 0xFF)
        LIGHT = RGBColor(0xCB, 0xD5, 0xE1)

        title_layout = prs.slide_layouts[0]   # Title slide
        content_layout = prs.slide_layouts[1]  # Title + content

        for i, slide_data in enumerate(slides):
            slide_title = slide_data.get("title", f"Slide {i+1}")
            slide_content = slide_data.get("content", "")

            layout = title_layout if i == 0 else content_layout
            slide = prs.slides.add_slide(layout)

            # Dark background
            bg = slide.background
            fill = bg.fill
            fill.solid()
            fill.fore_color.rgb = DARK_BG

            # Title
            if slide.shapes.title:
                tf = slide.shapes.title.text_frame
                tf.text = slide_title
                for para in tf.paragraphs:
                    for run in para.runs:
                        run.font.color.rgb = ACCENT
                        run.font.size = Pt(36 if i == 0 else 28)
                        run.font.bold = True

            # Content
            if slide_content and len(slide.placeholders) > 1:
                body = slide.placeholders[1]
                tf = body.text_frame
                tf.word_wrap = True
                tf.clear()

                for line in slide_content.split("\n"):
                    stripped = line.strip()
                    if not stripped:
                        continue
                    para = tf.add_paragraph()
                    if stripped.startswith("- ") or stripped.startswith("• "):
                        para.text = stripped[2:]
                        para.level = 1
                    else:
                        para.text = stripped
                        para.level = 0
                    for run in para.runs:
                        run.font.color.rgb = LIGHT
                        run.font.size = Pt(18)

        buf = io.BytesIO()
        prs.save(buf)
        filename = f"{title.replace(' ', '_')}_{datetime.utcnow().strftime('%Y%m%d%H%M%S')}.pptx"
        return {
            "file": _to_base64(buf),
            "filename": filename,
            "mime_type": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
            "success": True,
        }
    except ImportError:
        raise HTTPException(500, "python-pptx not installed")
    except Exception as e:
        logger.error(f"PPTX generation error: {e}")
        raise HTTPException(500, f"PPTX generation failed: {e}")


# ── HTML ──────────────────────────────────────────────────────────────────────

def generate_html(content: str, title: str = "Report") -> dict:
    import html as html_module

    # Convert markdown-ish content to HTML
    lines = content.split("\n")
    body_html = []
    in_ul = False

    for line in lines:
        stripped = line.strip()
        if not stripped:
            if in_ul:
                body_html.append("</ul>")
                in_ul = False
            body_html.append("<br>")
        elif stripped.startswith("## "):
            if in_ul:
                body_html.append("</ul>")
                in_ul = False
            body_html.append(f"<h2>{html_module.escape(stripped[3:])}</h2>")
        elif stripped.startswith("# "):
            if in_ul:
                body_html.append("</ul>")
                in_ul = False
            body_html.append(f"<h1>{html_module.escape(stripped[2:])}</h1>")
        elif stripped.startswith("### "):
            if in_ul:
                body_html.append("</ul>")
                in_ul = False
            body_html.append(f"<h3>{html_module.escape(stripped[4:])}</h3>")
        elif stripped.startswith("- ") or stripped.startswith("* "):
            if not in_ul:
                body_html.append("<ul>")
                in_ul = True
            body_html.append(f"<li>{html_module.escape(stripped[2:])}</li>")
        else:
            if in_ul:
                body_html.append("</ul>")
                in_ul = False
            body_html.append(f"<p>{html_module.escape(stripped)}</p>")

    if in_ul:
        body_html.append("</ul>")

    html_content = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <title>{html_module.escape(title)}</title>
  <style>
    *, *::before, *::after {{ box-sizing: border-box; margin: 0; padding: 0; }}
    body {{
      font-family: 'Segoe UI', system-ui, -apple-system, sans-serif;
      background: #0f172a;
      color: #e2e8f0;
      line-height: 1.7;
      padding: 2rem;
    }}
    .container {{ max-width: 900px; margin: 0 auto; }}
    header {{
      background: linear-gradient(135deg, #1e293b, #0f172a);
      border-left: 4px solid #38bdf8;
      padding: 1.5rem 2rem;
      margin-bottom: 2rem;
      border-radius: 0 8px 8px 0;
    }}
    header h1 {{ color: #38bdf8; font-size: 2rem; }}
    header small {{ color: #94a3b8; font-size: 0.85rem; }}
    .content {{ background: #1e293b; padding: 2rem; border-radius: 8px; }}
    h1, h2, h3 {{ color: #38bdf8; margin: 1.5rem 0 0.75rem; }}
    h1 {{ font-size: 1.8rem; }}
    h2 {{ font-size: 1.4rem; border-bottom: 1px solid #334155; padding-bottom: 0.4rem; }}
    h3 {{ font-size: 1.1rem; color: #7dd3fc; }}
    p {{ margin-bottom: 0.8rem; color: #cbd5e1; }}
    ul {{ padding-left: 1.5rem; margin-bottom: 0.8rem; }}
    li {{ margin-bottom: 0.4rem; color: #cbd5e1; }}
    footer {{
      text-align: center;
      margin-top: 2rem;
      color: #64748b;
      font-size: 0.8rem;
    }}
    @media (max-width: 600px) {{ body {{ padding: 1rem; }} }}
  </style>
</head>
<body>
  <div class="container">
    <header>
      <h1>{html_module.escape(title)}</h1>
      <small>Generated by S.T.E.W Agent • {datetime.utcnow().strftime('%Y-%m-%d %H:%M UTC')}</small>
    </header>
    <div class="content">
      {''.join(body_html)}
    </div>
    <footer>S.T.E.W — Structured Task Execution Workflow</footer>
  </div>
</body>
</html>"""

    buf = io.BytesIO(html_content.encode("utf-8"))
    filename = f"{title.replace(' ', '_')}_{datetime.utcnow().strftime('%Y%m%d%H%M%S')}.html"
    return {
        "file": _to_base64(buf),
        "filename": filename,
        "mime_type": "text/html",
        "success": True,
    }