gmal-scope-builder/backend/app/services/export_excel.py

"""Export ratecard data to Excel."""

import io
import logging
from collections import defaultdict

from openpyxl import Workbook
from openpyxl.styles import Font, Alignment, PatternFill, Border, Side
from openpyxl.utils import get_column_letter
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from app.models.gmal import GmalAsset, Role
from app.models.project import Project, ClientAsset, Match, RatecardLine

logger = logging.getLogger(__name__)

HEADER_FILL = PatternFill(start_color="1F4E79", end_color="1F4E79", fill_type="solid")
HEADER_FONT = Font(color="FFFFFF", bold=True, size=11)
DISCIPLINE_FILL = PatternFill(start_color="D6E4F0", end_color="D6E4F0", fill_type="solid")
THIN_BORDER = Border(
    left=Side(style="thin"),
    right=Side(style="thin"),
    top=Side(style="thin"),
    bottom=Side(style="thin"),
)


async def export_ratecard_excel(db: AsyncSession, project: Project) -> bytes:
    """Generate an Excel workbook with the ratecard data.

    Returns the workbook as bytes.
    """
    wb = Workbook()

    # Load all data
    lines_result = await db.execute(
        select(RatecardLine).where(RatecardLine.project_id == project.id)
    )
    lines = lines_result.scalars().all()

    if not lines:
        ws = wb.active
        ws.title = "Ratecard"
        ws["A1"] = "No ratecard data available"
        return _workbook_to_bytes(wb)

    # Load related entities
    role_ids = list(set(l.role_id for l in lines))
    asset_ids = list(set(l.client_asset_id for l in lines))
    gmal_ids = list(set(l.gmal_asset_id for l in lines))

    roles_result = await db.execute(select(Role).where(Role.id.in_(role_ids)))
    roles = {r.id: r for r in roles_result.scalars().all()}

    assets_result = await db.execute(select(ClientAsset).where(ClientAsset.id.in_(asset_ids)))
    client_assets = {a.id: a for a in assets_result.scalars().all()}

    gmals_result = await db.execute(select(GmalAsset).where(GmalAsset.id.in_(gmal_ids)))
    gmals = {g.id: g for g in gmals_result.scalars().all()}

    # Sheet 1: Ratecard Summary (roles x assets matrix)
    ws1 = wb.active
    ws1.title = "Ratecard Summary"
    _build_ratecard_sheet(ws1, lines, roles, client_assets, gmals)

    # Sheet 2: Asset Detail
    ws2 = wb.create_sheet("Asset Detail")
    await _build_asset_detail_sheet(ws2, db, project, client_assets, gmals)

    return _workbook_to_bytes(wb)


def _build_ratecard_sheet(ws, lines, roles, client_assets, gmals):
    """Build the main ratecard matrix: rows=roles, cols=client assets."""
    # Get unique sorted client assets and roles
    asset_ids_ordered = sorted(client_assets.keys())
    role_ids_ordered = sorted(roles.keys(), key=lambda rid: (roles[rid].discipline, roles[rid].sort_order or 0))

    # Build hours lookup: {(role_id, client_asset_id): total_hours}
    hours_map = {}
    for line in lines:
        effective_hours = line.manual_override if line.manual_override is not None else line.total_hours
        hours_map[(line.role_id, line.client_asset_id)] = float(effective_hours or 0)

    # Headers
    ws.cell(row=1, column=1, value="Discipline").font = HEADER_FONT
    ws.cell(row=1, column=1).fill = HEADER_FILL
    ws.cell(row=1, column=2, value="Role").font = HEADER_FONT
    ws.cell(row=1, column=2).fill = HEADER_FILL

    for col_idx, asset_id in enumerate(asset_ids_ordered, 3):
        ca = client_assets[asset_id]
        gmal_id = None
        for line in lines:
            if line.client_asset_id == asset_id:
                g = gmals.get(line.gmal_asset_id)
                gmal_id = g.gmal_id if g else None
                break

        header = f"{ca.raw_name}\n(Vol: {ca.volume})"
        if gmal_id:
            header += f"\n[{gmal_id}]"

        cell = ws.cell(row=1, column=col_idx, value=header)
        cell.font = HEADER_FONT
        cell.fill = HEADER_FILL
        cell.alignment = Alignment(wrap_text=True, horizontal="center")

    # Total column
    total_col = len(asset_ids_ordered) + 3
    ws.cell(row=1, column=total_col, value="Total Hours").font = HEADER_FONT
    ws.cell(row=1, column=total_col).fill = HEADER_FILL

    # Data rows
    current_discipline = None
    row_idx = 2

    for role_id in role_ids_ordered:
        role = roles[role_id]

        # Check if this role has any hours at all
        role_total = sum(hours_map.get((role_id, aid), 0) for aid in asset_ids_ordered)
        if role_total == 0:
            continue

        # Discipline grouping
        if role.discipline != current_discipline:
            current_discipline = role.discipline
            ws.cell(row=row_idx, column=1, value=current_discipline).font = Font(bold=True)
            ws.cell(row=row_idx, column=1).fill = DISCIPLINE_FILL
            for c in range(1, total_col + 1):
                ws.cell(row=row_idx, column=c).fill = DISCIPLINE_FILL
            row_idx += 1

        ws.cell(row=row_idx, column=1, value=role.discipline)
        ws.cell(row=row_idx, column=2, value=role.role_title)

        row_total = 0
        for col_idx, asset_id in enumerate(asset_ids_ordered, 3):
            hours = hours_map.get((role_id, asset_id), 0)
            if hours > 0:
                ws.cell(row=row_idx, column=col_idx, value=round(hours, 2))
                row_total += hours

        ws.cell(row=row_idx, column=total_col, value=round(row_total, 2)).font = Font(bold=True)
        row_idx += 1

    # Grand total row
    row_idx += 1
    ws.cell(row=row_idx, column=1, value="TOTAL").font = Font(bold=True, size=12)
    grand_total = 0
    for col_idx, asset_id in enumerate(asset_ids_ordered, 3):
        col_total = sum(hours_map.get((rid, asset_id), 0) for rid in role_ids_ordered)
        if col_total > 0:
            ws.cell(row=row_idx, column=col_idx, value=round(col_total, 2)).font = Font(bold=True)
            grand_total += col_total
    ws.cell(row=row_idx, column=total_col, value=round(grand_total, 2)).font = Font(bold=True, size=12)

    # Column widths
    ws.column_dimensions["A"].width = 25
    ws.column_dimensions["B"].width = 35
    for col_idx in range(3, total_col + 1):
        ws.column_dimensions[get_column_letter(col_idx)].width = 18


async def _build_asset_detail_sheet(ws, db, project, client_assets, gmals):
    """Build the asset detail sheet showing matches and caveats."""
    headers = ["Client Asset", "Volume", "Matched GMAL", "GMAL Name", "Confidence", "Score", "Caveats"]
    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

    # Load matches
    from app.models.project import Match
    matches_result = await db.execute(
        select(Match).where(
            Match.client_asset_id.in_(list(client_assets.keys())),
            Match.is_selected == True,
        )
    )
    matches = matches_result.scalars().all()
    match_by_asset = {m.client_asset_id: m for m in matches}

    row_idx = 2
    for asset_id in sorted(client_assets.keys()):
        ca = client_assets[asset_id]
        match = match_by_asset.get(asset_id)

        ws.cell(row=row_idx, column=1, value=ca.raw_name)
        ws.cell(row=row_idx, column=2, value=ca.volume)

        if match:
            gmal = gmals.get(match.gmal_asset_id)
            ws.cell(row=row_idx, column=3, value=gmal.gmal_id if gmal else "")
            ws.cell(row=row_idx, column=4, value=gmal.unique_name if gmal else "")
            ws.cell(row=row_idx, column=5, value=match.confidence.value)
            ws.cell(row=row_idx, column=6, value=float(match.confidence_score) if match.confidence_score else 0)
            ws.cell(row=row_idx, column=7, value=match.caveat_text or "")
        else:
            ws.cell(row=row_idx, column=3, value="No match")

        row_idx += 1

    # Column widths
    widths = [30, 10, 15, 40, 12, 10, 60]
    for i, w in enumerate(widths, 1):
        ws.column_dimensions[get_column_letter(i)].width = w


def _workbook_to_bytes(wb: Workbook) -> bytes:
    buf = io.BytesIO()
    wb.save(buf)
    buf.seek(0)
    return buf.read()