src/uofa_cli/excel_writer.py

"""Excel writer — writes LLM extraction results into a pack template with confidence coloring."""

from __future__ import annotations

import difflib
import shutil
from pathlib import Path

from uofa_cli.llm_extractor import ExtractionResult, FieldExtraction
from uofa_cli.excel_constants import (
    VV40_FACTOR_NAMES, NASA_ALL_FACTOR_NAMES,
    VV40_FACTOR_CATEGORIES, NASA_ONLY_FACTOR_CATEGORIES,
    VALID_DECISION_OUTCOMES, VALID_PROFILES, VALID_ASSURANCE_LEVELS,
    VALID_DEVICE_CLASSES,
    normalize_evidence_type,
)

# Confidence color thresholds
_GREEN = "C6EFCE"   # >= 0.85
_YELLOW = "FFEB9C"  # 0.50 - 0.84
_RED = "FFC7CE"     # < 0.50


def write_extraction(
    result: ExtractionResult,
    template_path: Path | None,
    output_path: Path,
    pack_name: str,
) -> Path:
    """Write extraction results to an Excel template.

    If template_path exists, copies it and fills in extracted values.
    If template_path is None, creates sheets from scratch.
    """
    import openpyxl
    from openpyxl.styles import PatternFill
    from openpyxl.comments import Comment

    if template_path and template_path.exists():
        shutil.copy2(template_path, output_path)
        wb = openpyxl.load_workbook(output_path)
    else:
        wb = openpyxl.Workbook()
        wb.remove(wb.active)
        _create_blank_sheets(wb, pack_name)

    # Ensure model_and_data has a Requirement entity. Real LLMs (especially
    # smaller local models) sometimes drop the Requirement when the COU is
    # a continuation of a previous one. Without this, the downstream
    # importer hard-fails with "Sheet Model & Data must have at least one
    # row with Entity Type = Requirement".
    _ensure_requirement_entity(result)

    # Write each section
    if "Assessment Summary" in wb.sheetnames:
        _write_summary_sheet(wb["Assessment Summary"], result.assessment_summary)

    if "Model & Data" in wb.sheetnames:
        _write_entities_sheet(wb["Model & Data"], result.model_and_data)

    if "Validation Results" in wb.sheetnames:
        _write_validation_sheet(wb["Validation Results"], result.validation_results)

    if "Credibility Factors" in wb.sheetnames:
        _write_factors_sheet(wb["Credibility Factors"], result.credibility_factors, pack_name)

    if "Decision" in wb.sheetnames:
        _write_decision_sheet(wb["Decision"], result.decision)

    wb.save(output_path)
    return output_path


def _create_blank_sheets(wb, pack_name: str) -> None:
    """Create blank template sheets when no template file exists."""
    # Assessment Summary
    ws = wb.create_sheet("Assessment Summary", 0)
    ws.append(["Unit of Assurance — Assessment Summary"])
    ws.append(["Project Name", "COU Name", "COU Description", "Profile",
               "Device Class", "Model Risk Level", "Assurance Level",
               "Standards Reference", "Assessor Name", "Assessment Date",
               "Source Document", "Has UQ?"])
    ws.append([None] * 12)  # Empty data row

    # Model & Data
    ws = wb.create_sheet("Model & Data")
    ws.append(["Model & Data"])
    ws.append(["Entity Type", "Name", "Identifier/URI", "Description", "Version", "Source"])

    # Validation Results
    ws = wb.create_sheet("Validation Results")
    ws.append(["Validation Results"])
    ws.append(["Result Name", "Type", "Identifier/URI", "Description",
               "Compares To", "Has UQ?", "UQ Method", "Metric Value", "Pass/Fail"])

    # Credibility Factors
    ws = wb.create_sheet("Credibility Factors")
    ws.append(["Credibility Factors"])
    ws.append([""])
    ws.append(["Factor Type", "Category", "Required Level", "Achieved Level",
               "Acceptance Criteria", "Rationale", "Factor Status", "Linked Evidence"])
    ws.append(["(pre-filled)", "(pre-filled)", "(1-5)", "(1-5)",
               "(optional)", "(optional)", "(status)", "(URI)"])

    # Pre-fill factor rows
    is_nasa = "nasa" in pack_name.lower()
    categories = list(VV40_FACTOR_CATEGORIES)
    if is_nasa:
        categories.extend(NASA_ONLY_FACTOR_CATEGORIES)
    for name, cat in categories:
        ws.append([name, cat, None, None, None, None, None, None])

    # Decision
    ws = wb.create_sheet("Decision")
    ws.append(["Decision"])
    ws.append(["Decision Outcome", "Decision Rationale", "Criteria Set",
               "Decided By", "Decision Date"])
    ws.append([None] * 5)


# ── Sheet writers ────────────────────────────────────────────


def _write_summary_sheet(ws, summary: dict[str, FieldExtraction]) -> None:
    """Write assessment summary fields to the data row."""
    # Find the data row (row after header row)
    data_row = _find_data_row(ws)

    field_map = {
        "project_name": 1,
        "cou_name": 2,
        "cou_description": 3,
        "profile": 4,
        "device_class": 5,
        "model_risk_level": 6,
        "assurance_level": 7,
        "standards_reference": 8,
        "assessor_name": 9,
        "has_uq": 12,
    }

    dropdown_maps = {
        "profile": VALID_PROFILES,
        "device_class": VALID_DEVICE_CLASSES,
        "assurance_level": VALID_ASSURANCE_LEVELS,
    }

    for field_name, col in field_map.items():
        fe = summary.get(field_name)
        cell = ws.cell(row=data_row, column=col)
        if fe is None or fe.value is None:
            # Clear the template's placeholder hint text so it doesn't
            # surface to the importer as if it were extracted data. Without
            # this, a missing extraction looks like "Profile = 'Minimal or
            # Complete'" downstream — a confusing enum-validation error
            # instead of a clear "required field empty" signal.
            cell.value = None
            continue

        value = fe.value
        # Fuzzy match dropdowns
        if field_name in dropdown_maps:
            value = _fuzzy_match_dropdown(str(value), dropdown_maps[field_name])

        cell.value = value
        _apply_confidence_color(cell, fe.confidence)
        _add_source_comment(cell, fe)


def _write_entities_sheet(ws, entities: list[dict[str, FieldExtraction]]) -> None:
    """Write model & data entities."""
    data_row = _find_data_row(ws)

    for i, entity in enumerate(entities):
        row = data_row + i
        field_map = {
            "entity_type": 1,
            "name": 2,
            "description": 4,
        }
        for field_name, col in field_map.items():
            fe = entity.get(field_name)
            if fe is None or fe.value is None:
                continue
            cell = ws.cell(row=row, column=col)
            cell.value = fe.value
            _apply_confidence_color(cell, fe.confidence)
            _add_source_comment(cell, fe)


def _write_validation_sheet(ws, results: list[dict[str, FieldExtraction]]) -> None:
    """Write validation result rows."""
    from openpyxl.comments import Comment
    from openpyxl.styles import PatternFill

    data_row = _find_data_row(ws)

    for i, vr in enumerate(results):
        row = data_row + i
        field_map = {
            "name": 1,
            "evidence_type": 2,
            "description": 4,
            "compares_to": 5,
            "has_uq": 6,
            "uq_method": 7,
            "metric_value": 8,
            "pass_fail": 9,
        }
        for field_name, col in field_map.items():
            fe = vr.get(field_name)
            if fe is None or fe.value is None:
                continue

            value = fe.value
            substitution_note = None

            # Normalize evidence_type against the canonical enum so the
            # downstream importer accepts what the LLM produced. Without this,
            # descriptive labels like 'GridConvergenceStudy' would land in the
            # cell verbatim and the importer would reject them.
            if field_name == "evidence_type":
                original = str(value)
                normalized, substituted = normalize_evidence_type(original)
                if substituted:
                    substitution_note = (
                        f"LLM emitted '{original}'; normalized to "
                        f"'{normalized}' to match the canonical evidence_type enum."
                    )
                value = normalized

            cell = ws.cell(row=row, column=col)
            cell.value = value
            _apply_confidence_color(cell, fe.confidence)
            if substitution_note:
                # Yellow-highlight the substituted cell so a human reviewer sees
                # it before saving / re-running. Comment carries the original.
                cell.fill = PatternFill(start_color=_YELLOW, end_color=_YELLOW, fill_type="solid")
                cell.comment = Comment(substitution_note, "uofa-extract")
            _add_source_comment(cell, fe)


def _write_factors_sheet(
    ws, factors: list[dict[str, FieldExtraction]], pack_name: str,
) -> None:
    """Write credibility factor values to the existing factor rows in the template.

    Matches extracted factors to existing rows by factor_type in column A.
    Column H (Linked Evidence) is explicitly left empty.
    """
    is_nasa = "nasa" in pack_name.lower()
    valid_names = NASA_ALL_FACTOR_NAMES if is_nasa else VV40_FACTOR_NAMES

    # Build a map of factor_type -> row number from the existing sheet
    factor_rows: dict[str, int] = {}
    for row in range(1, ws.max_row + 1):
        cell_val = ws.cell(row=row, column=1).value
        if cell_val and str(cell_val) in valid_names:
            factor_rows[str(cell_val)] = row

    # If no pre-existing rows (blank template without factors), start after headers
    if not factor_rows:
        start_row = 5  # After header rows
        for i, factor in enumerate(factors):
            ft_fe = factor.get("factor_type")
            if ft_fe is None:
                continue
            row = start_row + i
            ws.cell(row=row, column=1).value = ft_fe.value
            _write_factor_row(ws, row, factor)
        return

    # Match extracted factors to existing rows
    for factor in factors:
        ft_fe = factor.get("factor_type")
        if ft_fe is None or ft_fe.value is None:
            continue

        ft_name = str(ft_fe.value)
        row = factor_rows.get(ft_name)
        if row is None:
            # Try fuzzy match
            matches = difflib.get_close_matches(ft_name, list(factor_rows.keys()), n=1, cutoff=0.6)
            if matches:
                row = factor_rows[matches[0]]

        if row is not None:
            _write_factor_row(ws, row, factor)


def _write_factor_row(ws, row: int, factor: dict[str, FieldExtraction]) -> None:
    """Write a single factor's data to a row. Column H (8) is left empty."""
    field_map = {
        "required_level": 3,     # Column C
        "achieved_level": 4,     # Column D
        "acceptance_criteria": 5, # Column E
        "rationale": 6,          # Column F
        "status": 7,             # Column G
        # Column H (8) = Linked Evidence — intentionally skipped
    }

    for field_name, col in field_map.items():
        fe = factor.get(field_name)
        if fe is None or fe.value is None:
            continue
        cell = ws.cell(row=row, column=col)
        cell.value = fe.value
        _apply_confidence_color(cell, fe.confidence)
        _add_source_comment(cell, fe)


def _write_decision_sheet(ws, decision: dict[str, FieldExtraction]) -> None:
    """Write decision fields."""
    data_row = _find_data_row(ws)

    field_map = {
        "outcome": 1,
        "rationale": 2,
        "decided_by": 4,
        "decision_date": 5,
    }

    for field_name, col in field_map.items():
        fe = decision.get(field_name)
        cell = ws.cell(row=data_row, column=col)
        if fe is None or fe.value is None:
            # Clear template placeholder hint (see _write_summary_sheet
            # for the rationale).
            cell.value = None
            continue

        value = fe.value
        if field_name == "outcome":
            value = _fuzzy_match_dropdown(str(value), VALID_DECISION_OUTCOMES)

        cell.value = value
        _apply_confidence_color(cell, fe.confidence)
        _add_source_comment(cell, fe)


# ── Helpers ──────────────────────────────────────────────────


def _find_data_row(ws) -> int:
    """Find the first data row (row after the header row).

    Looks for a row with header-like content, returns the row after it.
    Falls back to row 3 (typical layout: row 1 = title, row 2 = headers, row 3 = data).
    """
    for row in range(1, min(ws.max_row + 1, 10)):
        cell = ws.cell(row=row, column=1).value
        if cell and any(kw in str(cell) for kw in [
            "Project Name", "Entity Type", "Result Name",
            "Decision Outcome", "Factor Type",
        ]):
            return row + 1
    return 3


def _apply_confidence_color(cell, confidence: float) -> None:
    """Apply green/yellow/red fill based on confidence score."""
    from openpyxl.styles import PatternFill

    if confidence >= 0.85:
        fill_color = _GREEN
    elif confidence >= 0.50:
        fill_color = _YELLOW
    else:
        fill_color = _RED

    cell.fill = PatternFill(start_color=fill_color, end_color=fill_color, fill_type="solid")


def _add_source_comment(cell, fe: FieldExtraction) -> None:
    """Add a cell comment with source attribution."""
    from openpyxl.comments import Comment

    parts = ["[Extract]"]
    if fe.source_file:
        source = f"Source: {fe.source_file}"
        if fe.source_page:
            source += f" p.{fe.source_page}"
        parts.append(source)
    parts.append(f"Confidence: {fe.confidence:.0%}")

    cell.comment = Comment(" | ".join(parts), "uofa extract")


def _fuzzy_match_dropdown(value: str, valid_values: list[str]) -> str:
    """Fuzzy match a value against valid dropdown options.

    Handles three classes of LLM-output gap on enum fields:

    1. Exact match — return as-is.
    2. Enumerator-text echo: LLM emitted prompt text like ``"Minimal or
       Complete"``, ``"Low / Medium / High"``, or ``"Accepted / Not
       accepted / Conditional"`` instead of picking one option. Split on
       common separators (``or``, ``/``, ``,``) and return the first token
       that matches a valid value.
    3. Typos / minor variants: difflib fuzzy match (cutoff 0.6) — handles
       things like ``"ValidationResults"`` vs ``"ValidationResult"``.

    Falls back to the original value if nothing matches.
    """
    if value in valid_values:
        return value

    import re
    tokens = re.split(r"\s+or\s+|\s*[/,]\s*", value)
    for tok in tokens:
        tok = tok.strip()
        if tok in valid_values:
            return tok

    matches = difflib.get_close_matches(value, valid_values, n=1, cutoff=0.6)
    return matches[0] if matches else value


def _ensure_requirement_entity(result: ExtractionResult) -> None:
    """Synthesize a Requirement entity from the Assessment Summary when
    ``result.model_and_data`` has none. Mutates the list in place.

    Real LLMs (especially smaller local models) sometimes drop the
    Requirement when the COU is a continuation of a previous one (the
    model decides "this reuses cou1's setup" and emits only Model +
    Dataset). The downstream importer hard-fails without a Requirement,
    so we backfill from the COU's own description in the Assessment
    Summary. The synthesized row carries confidence 0.5 (yellow band)
    so a human reviewer notices it before signing.
    """
    has_requirement = any(
        e.get("entity_type") and getattr(e["entity_type"], "value", None) == "Requirement"
        for e in result.model_and_data
    )
    if has_requirement:
        return

    summary = result.assessment_summary or {}
    cou_name_fe = summary.get("cou_name")
    cou_desc_fe = summary.get("cou_description")
    cou_name = cou_name_fe.value if cou_name_fe and cou_name_fe.value else None
    cou_desc = cou_desc_fe.value if cou_desc_fe and cou_desc_fe.value else None

    synth_name = cou_name or "Implicit COU Requirement"
    synth_desc = (
        f"Auto-synthesized from COU context: {cou_desc or cou_name or 'unspecified'}. "
        "The extractor did not emit a Requirement entity. Review and replace "
        "with the explicit top-level performance, safety, or certification "
        "requirement before signing."
    )

    result.model_and_data.insert(0, {
        "entity_type": FieldExtraction(
            value="Requirement", confidence=0.5, source_file="auto-synthesized"),
        "name": FieldExtraction(
            value=synth_name, confidence=0.5, source_file="auto-synthesized"),
        "description": FieldExtraction(
            value=synth_desc, confidence=0.5, source_file="auto-synthesized"),
    })