personal-ai-assistant/backend/app/services/pdf_service.py

import uuid
from datetime import datetime, timezone
from pathlib import Path

from jinja2 import Environment, FileSystemLoader, select_autoescape
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.models.document import Document
from app.models.generated_pdf import GeneratedPdf
from app.models.user import User
from app.services.memory_service import get_user_memories

TEMPLATE_DIR = Path(__file__).parent.parent / "templates" / "pdf"
_file_jinja_env = Environment(
    loader=FileSystemLoader(str(TEMPLATE_DIR)),
    autoescape=select_autoescape(["html"]),
)
_string_jinja_env = Environment(autoescape=select_autoescape(["html"]))


async def generate_pdf_report(
    db: AsyncSession,
    user_id: uuid.UUID,
    title: str,
    document_ids: list[uuid.UUID] | None = None,
    chat_id: uuid.UUID | None = None,
    template_id: uuid.UUID | None = None,
) -> GeneratedPdf:
    # Load user
    result = await db.execute(select(User).where(User.id == user_id))
    user = result.scalar_one()

    # Load memories
    memories = await get_user_memories(db, user_id, is_active=True)
    memory_data = [
        {"category": m.category, "title": m.title, "content": m.content, "importance": m.importance}
        for m in memories
    ]

    # Load documents
    doc_data = []
    if document_ids:
        for doc_id in document_ids:
            result = await db.execute(
                select(Document).where(Document.id == doc_id, Document.user_id == user_id)
            )
            doc = result.scalar_one_or_none()
            if doc:
                doc_data.append({
                    "original_filename": doc.original_filename,
                    "doc_type": doc.doc_type,
                    "excerpt": (doc.extracted_text or "")[:2000],
                })

    template_vars = dict(
        title=title,
        user_name=user.full_name or user.username,
        generated_at=datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC"),
        memories=memory_data,
        documents=doc_data,
        ai_summary=None,
    )

    # Resolve template
    resolved_template_id = template_id
    if template_id:
        from app.services.pdf_template_service import get_template
        tmpl = await get_template(db, template_id)
        jinja_template = _string_jinja_env.from_string(tmpl.html_content)
    else:
        from app.services.pdf_template_service import get_default_template
        default_tmpl = await get_default_template(db)
        if default_tmpl:
            resolved_template_id = default_tmpl.id
            jinja_template = _string_jinja_env.from_string(default_tmpl.html_content)
        else:
            jinja_template = _file_jinja_env.get_template("report.html")

    html = jinja_template.render(**template_vars)

    # Generate PDF
    pdf_id = uuid.uuid4()
    pdf_dir = Path(settings.UPLOAD_DIR).parent / "pdfs" / str(user_id)
    pdf_dir.mkdir(parents=True, exist_ok=True)
    pdf_path = pdf_dir / f"{pdf_id}.pdf"

    try:
        from weasyprint import HTML
        HTML(string=html).write_pdf(str(pdf_path))
    except ImportError:
        pdf_path = pdf_path.with_suffix(".html")
        pdf_path.write_text(html, encoding="utf-8")

    # Save record
    generated = GeneratedPdf(
        id=pdf_id,
        user_id=user_id,
        title=title,
        storage_path=str(pdf_path),
        source_document_ids=document_ids,
        source_chat_id=chat_id,
        template_id=resolved_template_id,
    )
    db.add(generated)
    await db.flush()
    await db.refresh(generated)
    return generated


async def get_user_pdfs(db: AsyncSession, user_id: uuid.UUID) -> list[GeneratedPdf]:
    result = await db.execute(
        select(GeneratedPdf).where(GeneratedPdf.user_id == user_id)
        .order_by(GeneratedPdf.created_at.desc())
    )
    return list(result.scalars().all())


async def get_pdf(db: AsyncSession, pdf_id: uuid.UUID, user_id: uuid.UUID) -> GeneratedPdf | None:
    result = await db.execute(
        select(GeneratedPdf).where(GeneratedPdf.id == pdf_id, GeneratedPdf.user_id == user_id)
    )
    return result.scalar_one_or_none()