HC900-Crawler/mcp-server/instrument_inference/excel.py

"""3시트 Excel 초안 생성기 (§3 스키마)."""
from __future__ import annotations
import os
import tempfile
from datetime import datetime

import openpyxl
from openpyxl.styles import Font, PatternFill, Alignment

from .rules import get_all_measurements, get_all_modifiers, get_all_special_prefixes


INSTRUMENT_COLS = [
    "instrument_id", "display_name", "parent_base_tag", "role", "loop",
    "area", "measures", "data_points", "from", "to", "description",
    "confidence", "needs_review", "inference_basis", "operator_notes", "delete",
]

POWER_EQUIPMENT_COLS = [
    "instrument_id", "display_name", "parent_base_tag", "role", "loop",
    "area", "equipment_type", "data_points", "description",
    "confidence", "needs_review", "inference_basis", "operator_notes", "delete",
]

HEADER_FILL = PatternFill(start_color="4472C4", end_color="4472C4", fill_type="solid")
HEADER_FONT = Font(bold=True, color="FFFFFF", size=11)
LOW_FILL = PatternFill(start_color="FFC7CE", end_color="FFC7CE", fill_type="solid")
MEDIUM_FILL = PatternFill(start_color="FFEB9C", end_color="FFEB9C", fill_type="solid")


def generate_excel(instruments: list[dict], unmatched: list[str], power_equipment: list[dict] | None = None) -> str:
    """
    4시트 Excel 생성 후 파일 경로 반환.

    Args:
        instruments:     infer.py 결과 (instrument dict 리스트)
        unmatched:       룰 미매칭 base_tag 목록
        power_equipment: 동력기기 (펌프, 압축기, 교반기 등)
        output_dir:      출력 디렉토리 (기본: /tmp)

    Returns:
        생성된 xlsx 파일의 절대 경로
    """
    output_dir = tempfile.gettempdir()
    os.makedirs(output_dir, exist_ok=True)

    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    filepath = os.path.join(output_dir, f"instruments_draft_{ts}.xlsx")

    wb = openpyxl.Workbook()

    # ── 시트 1: instruments ──
    _build_instruments_sheet(wb, instruments)

    # ── 시트 2: power_equipment ──
    _build_power_equipment_sheet(wb, power_equipment or [])

    # ── 시트 3: unmatched_tags ──
    _build_unmatched_sheet(wb, unmatched)

    # ── 시트 4: naming_convention ──
    _build_naming_convention_sheet(wb)

    wb.save(filepath)
    return filepath


def _build_instruments_sheet(wb: openpyxl.Workbook, instruments: list[dict]) -> None:
    ws = wb.active
    ws.title = "instruments"

    # 헤더
    for col_idx, col_name in enumerate(INSTRUMENT_COLS, 1):
        cell = ws.cell(row=1, column=col_idx, value=col_name)
        cell.font = HEADER_FONT
        cell.fill = HEADER_FILL
        cell.alignment = Alignment(horizontal="center", wrap_text=True)

    # 데이터 행
    for row_idx, inst in enumerate(instruments, 2):
        for col_idx, col_name in enumerate(INSTRUMENT_COLS, 1):
            val = inst.get(col_name, "")
            cell = ws.cell(row=row_idx, column=col_idx, value=val)

            # confidence별 색상
            if col_name == "confidence":
                if val == "low":
                    cell.fill = LOW_FILL
                elif val == "medium":
                    cell.fill = MEDIUM_FILL

    # 컬럼 너비 자동 조정
    for col_idx in range(1, len(INSTRUMENT_COLS) + 1):
        max_len = 0
        for row in range(1, ws.max_row + 1):
            val = str(ws.cell(row=row, column=col_idx).value or "")
            max_len = max(max_len, len(val))
        ws.column_dimensions[openpyxl.utils.get_column_letter(col_idx)].width = min(max_len + 4, 40)


def _build_naming_convention_sheet(wb: openpyxl.Workbook) -> None:
    ws = wb.create_sheet("naming_convention")

    # 측정량 표
    ws.cell(row=1, column=1, value="측정량 (첫 글자)").font = HEADER_FONT
    ws.cell(row=1, column=2, value="의미").font = HEADER_FONT
    meas = get_all_measurements()
    for i, (letter, meaning) in enumerate(meas.items(), 2):
        ws.cell(row=i, column=1, value=letter)
        ws.cell(row=i, column=2, value=meaning)

    # 수식어 표 (D열부터)
    ws.cell(row=1, column=4, value="수식어 (두 번째 이후)").font = HEADER_FONT
    ws.cell(row=1, column=5, value="role").font = HEADER_FONT
    ws.cell(row=1, column=6, value="가상").font = HEADER_FONT
    mods = get_all_modifiers()
    for i, (letter, info) in enumerate(mods.items(), 2):
        ws.cell(row=i, column=4, value=letter)
        ws.cell(row=i, column=5, value=info.get("role", ""))
        ws.cell(row=i, column=6, value=info.get("virtual", False))

    # 특수 prefix 표 (G열부터)
    ws.cell(row=1, column=7, value="특수 prefix").font = HEADER_FONT
    ws.cell(row=1, column=8, value="role").font = HEADER_FONT
    sps = get_all_special_prefixes()
    for i, (prefix, info) in enumerate(sps.items(), 2):
        ws.cell(row=i, column=7, value=prefix)
        ws.cell(row=i, column=8, value=info.get("role", ""))


def _build_unmatched_sheet(wb: openpyxl.Workbook, unmatched: list[str]) -> None:
    ws = wb.create_sheet("unmatched_tags")

    ws.cell(row=1, column=1, value="base_tag").font = HEADER_FONT
    ws.cell(row=1, column=2, value="area").font = HEADER_FONT
    ws.cell(row=1, column=3, value="action").font = HEADER_FONT
    ws.cell(row=1, column=4, value="operator_notes").font = HEADER_FONT

    for i, tag in enumerate(unmatched, 2):
        ws.cell(row=i, column=1, value=tag)
        ws.cell(row=i, column=2, value="(none)")
        ws.cell(row=i, column=3, value="운영자가 instruments 시트에 행 추가 필요")
        ws.cell(row=i, column=4, value="")


def _build_power_equipment_sheet(wb: openpyxl.Workbook, equipment: list[dict]) -> None:
    ws = wb.create_sheet("power_equipment")

    for col_idx, col_name in enumerate(POWER_EQUIPMENT_COLS, 1):
        cell = ws.cell(row=1, column=col_idx, value=col_name)
        cell.font = HEADER_FONT
        cell.fill = HEADER_FILL
        cell.alignment = Alignment(horizontal="center", wrap_text=True)

    for row_idx, eq in enumerate(equipment, 2):
        for col_idx, col_name in enumerate(POWER_EQUIPMENT_COLS, 1):
            val = eq.get(col_name, "")
            cell = ws.cell(row=row_idx, column=col_idx, value=val)

            if col_name == "confidence":
                if val == "low":
                    cell.fill = LOW_FILL
                elif val == "medium":
                    cell.fill = MEDIUM_FILL

    for col_idx in range(1, len(POWER_EQUIPMENT_COLS) + 1):
        max_len = 0
        for row in range(1, ws.max_row + 1):
            val = str(ws.cell(row=row, column=col_idx).value or "")
            max_len = max(max_len, len(val))
        ws.column_dimensions[openpyxl.utils.get_column_letter(col_idx)].width = min(max_len + 4, 40)