src/totem_workbook.py

from __future__ import annotations

from pathlib import Path
from tempfile import NamedTemporaryFile
from typing import Any, Union
import math

import pandas as pd
from openpyxl import load_workbook


APP_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_WORKBOOK = APP_ROOT / "data" / "order69_macmillan_totem_rebuilt.xlsx"

METRICS = [
    "Clarity",
    "Rhythm",
    "Read-aloud Flow",
    "Emotional Truth",
    "Visual Strength",
    "Commercial Publishability",
]

LOG_COLUMNS = [
    "Sequence",
    "Stanza ID",
    "Draft / Pass",
    *METRICS,
    "Weighted Score",
    "Average",
    "Gate",
    "Revision Flag",
    "Priority Fix",
    "Notes",
]

KEY_READ_SHEETS = [
    "IDENTITY",
    "CANON",
    "VALUES",
    "STORY",
    "PITCH",
    "BRAND",
    "TONE",
    "SATIRE",
    "HANDOFF",
    "RECENT_CONTEXT",
    "CHAR_HENRY",
]

UPLOAD_TYPES = Union[str, Path, Any]


def workbook_path(uploaded_file: UPLOAD_TYPES | None = None) -> Path:
    if uploaded_file is None:
        return DEFAULT_WORKBOOK
    if isinstance(uploaded_file, (str, Path)):
        return Path(uploaded_file)
    if hasattr(uploaded_file, "name"):
        return Path(uploaded_file.name)
    return DEFAULT_WORKBOOK


def _text(value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, float) and math.isnan(value):
        return ""
    return str(value).strip()


def _number(value: Any) -> float | None:
    if value is None or value == "":
        return None
    if isinstance(value, float) and math.isnan(value):
        return None
    if isinstance(value, str) and value.startswith("="):
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _load(path: Path, data_only: bool = False):
    return load_workbook(path, data_only=data_only, read_only=False, keep_vba=path.suffix.lower() == ".xlsm")


def table_from_sheet(path: Path, sheet_name: str, header_row: int, start_row: int | None = None) -> pd.DataFrame:
    wb = _load(path)
    ws = wb[sheet_name]
    start = start_row or header_row + 1
    headers = [_text(ws.cell(header_row, col).value) for col in range(1, ws.max_column + 1)]
    rows: list[list[str]] = []

    for row_index in range(start, ws.max_row + 1):
        row = [_text(ws.cell(row_index, col).value) for col in range(1, ws.max_column + 1)]
        if any(row):
            rows.append(row)

    width = max(len(headers), max((len(row) for row in rows), default=0))
    headers = (headers + [f"Column {idx}" for idx in range(len(headers) + 1, width + 1)])[:width]
    normalized = [(row + [""] * width)[:width] for row in rows]
    df = pd.DataFrame(normalized, columns=headers)
    return df.loc[:, [col for col in df.columns if col]]


def workbook_overview(path: Path) -> dict[str, Any]:
    wb = _load(path)
    sheets = []
    for ws in wb.worksheets:
        nonempty = sum(1 for cell in ws._cells.values() if cell.value not in (None, ""))
        sheets.append(
            {
                "Sheet": ws.title,
                "Rows": ws.max_row,
                "Columns": ws.max_column,
                "Filled cells": nonempty,
            }
        )

    chain = table_from_sheet(path, "Chain", 2)
    top_roles = chain.head(8).to_dict("records") if not chain.empty else []
    return {
        "sheet_count": len(wb.sheetnames),
        "filled_cells": sum(row["Filled cells"] for row in sheets),
        "sheets": pd.DataFrame(sheets),
        "top_roles": top_roles,
    }


def chain_table(path: Path) -> pd.DataFrame:
    return table_from_sheet(path, "Chain", 2)


def protocol_table(path: Path) -> pd.DataFrame:
    df = table_from_sheet(path, "TOTEM_PROTOCOL", 5)
    if "Metric" in df.columns:
        df = df[df["Metric"].isin(METRICS)].copy()
    if "Weight" in df.columns:
        df["Weight"] = pd.to_numeric(df["Weight"], errors="coerce")
    return df


def protocol_weights(path: Path) -> dict[str, float]:
    df = protocol_table(path)
    weights = {row["Metric"]: float(row["Weight"]) for _, row in df.iterrows() if row.get("Metric") in METRICS}
    if not weights:
        weights = {
            "Clarity": 0.20,
            "Rhythm": 0.15,
            "Read-aloud Flow": 0.20,
            "Emotional Truth": 0.15,
            "Visual Strength": 0.15,
            "Commercial Publishability": 0.15,
        }
    return weights


def gate_for_scores(scores: dict[str, float], weights: dict[str, float]) -> dict[str, Any]:
    clean_scores = {metric: _number(scores.get(metric)) for metric in METRICS}
    present = {metric: score for metric, score in clean_scores.items() if score is not None}

    if not present:
        return {
            "Weighted Score": "",
            "Average": "",
            "Gate": "",
            "Revision Flag": "",
            "Priority Fix": "",
        }

    weighted = round(sum(float(present.get(metric, 0)) * weights.get(metric, 0) for metric in METRICS), 1)
    average = round(sum(present.values()) / len(present), 1)
    lowest_metric = min(present, key=lambda metric: present[metric])
    lowest_score = present[lowest_metric]
    low_count = sum(1 for score in present.values() if score <= 6)

    rhythm = present.get("Rhythm")
    flow = present.get("Read-aloud Flow")
    commercial = present.get("Commercial Publishability")

    if lowest_score <= 4:
        gate = "HARD FAIL"
    elif low_count >= 2:
        gate = "SOFT FAIL"
    elif (rhythm is not None and rhythm < 7) or (flow is not None and flow < 7):
        gate = "READ-ALOUD BLOCK"
    elif commercial is not None and commercial < 7:
        gate = "COMMERCIAL CHECK"
    elif weighted >= 8 and lowest_score >= 7:
        gate = "GREENLIGHT"
    else:
        gate = "REVISE"

    return {
        "Weighted Score": weighted,
        "Average": average,
        "Gate": gate,
        "Revision Flag": "No" if gate == "GREENLIGHT" else "Yes",
        "Priority Fix": lowest_metric,
    }


def score_log(path: Path) -> pd.DataFrame:
    wb = _load(path, data_only=False)
    ws = wb["TOTEM_LOG"]
    weights = protocol_weights(path)
    rows: list[dict[str, Any]] = []

    for row_index in range(7, min(ws.max_row, 86) + 1):
        raw = {
            "Sequence": _text(ws.cell(row_index, 1).value),
            "Stanza ID": _text(ws.cell(row_index, 2).value),
            "Draft / Pass": _text(ws.cell(row_index, 3).value),
            "Clarity": _number(ws.cell(row_index, 4).value),
            "Rhythm": _number(ws.cell(row_index, 5).value),
            "Read-aloud Flow": _number(ws.cell(row_index, 6).value),
            "Emotional Truth": _number(ws.cell(row_index, 7).value),
            "Visual Strength": _number(ws.cell(row_index, 8).value),
            "Commercial Publishability": _number(ws.cell(row_index, 9).value),
            "Priority Fix": _text(ws.cell(row_index, 14).value),
            "Notes": _text(ws.cell(row_index, 15).value),
        }
        priority_cell = raw["Priority Fix"]
        priority_is_formula = priority_cell.startswith("=")
        has_user_content = any(raw.get(col) not in ("", None) for col in ["Sequence", "Stanza ID", "Draft / Pass", *METRICS, "Notes"])
        has_user_content = has_user_content or bool(priority_cell and not priority_is_formula)
        if not has_user_content:
            continue

        calculated = gate_for_scores({metric: raw[metric] for metric in METRICS}, weights)
        if raw["Priority Fix"] and raw["Priority Fix"] not in METRICS and not priority_is_formula:
            raw["Notes"] = raw["Notes"] or raw["Priority Fix"]
            raw["Priority Fix"] = calculated["Priority Fix"]
        elif not raw["Priority Fix"] or priority_is_formula:
            raw["Priority Fix"] = calculated["Priority Fix"]

        raw.update(
            {
                "Weighted Score": calculated["Weighted Score"],
                "Average": calculated["Average"],
                "Gate": calculated["Gate"],
                "Revision Flag": calculated["Revision Flag"],
            }
        )
        rows.append(raw)

    return pd.DataFrame(rows, columns=LOG_COLUMNS)


def recalculate_log(log_df: pd.DataFrame | None, path: Path) -> pd.DataFrame:
    if log_df is None or log_df.empty:
        return pd.DataFrame(columns=LOG_COLUMNS)

    weights = protocol_weights(path)
    rows: list[dict[str, Any]] = []
    for _, row in log_df.iterrows():
        item = {column: row.get(column, "") for column in LOG_COLUMNS}
        scores = {metric: _number(item.get(metric)) for metric in METRICS}
        has_content = any(_text(item.get(col)) for col in ["Sequence", "Stanza ID", "Draft / Pass", "Priority Fix", "Notes"]) or any(
            value is not None for value in scores.values()
        )
        if not has_content:
            continue
        calculated = gate_for_scores(scores, weights)
        item.update(calculated)
        rows.append(item)

    return pd.DataFrame(rows, columns=LOG_COLUMNS)


def score_single_row(
    path: Path,
    sequence: str,
    stanza_id: str,
    draft_pass: str,
    clarity: float,
    rhythm: float,
    flow: float,
    emotional_truth: float,
    visual_strength: float,
    commercial: float,
    notes: str,
) -> pd.DataFrame:
    scores = {
        "Clarity": clarity,
        "Rhythm": rhythm,
        "Read-aloud Flow": flow,
        "Emotional Truth": emotional_truth,
        "Visual Strength": visual_strength,
        "Commercial Publishability": commercial,
    }
    calculated = gate_for_scores(scores, protocol_weights(path))
    row = {
        "Sequence": sequence,
        "Stanza ID": stanza_id,
        "Draft / Pass": draft_pass,
        **scores,
        **calculated,
        "Notes": notes,
    }
    return pd.DataFrame([row], columns=LOG_COLUMNS)


def viability_table(path: Path) -> tuple[pd.DataFrame, str]:
    wb = _load(path, data_only=False)
    ws = wb["VIABILITY"]
    rows = []
    for row_index in range(5, ws.max_row + 1):
        metric = _text(ws.cell(row_index, 1).value)
        score = _number(ws.cell(row_index, 2).value)
        read = _text(ws.cell(row_index, 3).value)
        if metric and score is not None:
            rows.append({"Metric": metric, "Score": score, "Read": read})
    df = pd.DataFrame(rows)
    if df.empty:
        return df, "No viability rows found."

    avg = round(float(df["Score"].mean()), 1)
    strong = int((df["Score"] >= 8).sum())
    needs_work = int((df["Score"] < 7).sum())
    weakest = df.loc[df["Score"].idxmin()]
    summary = (
        f"Average viability: {avg}/10. Strong metrics: {strong}. "
        f"Needs work under 7: {needs_work}. Weakest commercial pressure point: "
        f"{weakest['Metric']} ({weakest['Score']}/10)."
    )
    return df, summary


def workstack_table(path: Path) -> pd.DataFrame:
    return table_from_sheet(path, "WORKSTACK", 2)


def manuscript_tracker_table(path: Path) -> pd.DataFrame:
    return table_from_sheet(path, "MANUSCRIPT_TRACKER", 4)


def command_registry_table(path: Path) -> pd.DataFrame:
    return table_from_sheet(path, "COMMAND_REGISTRY", 4)


def key_reads_markdown(path: Path) -> str:
    wb = _load(path)
    chunks = []
    for sheet_name in KEY_READ_SHEETS:
        if sheet_name not in wb.sheetnames:
            continue
        ws = wb[sheet_name]
        title = _text(ws["A1"].value) or sheet_name
        purpose = _text(ws["B2"].value)
        current = _text(ws["B3"].value)
        note = _text(ws["B4"].value)
        body = current or purpose or note
        if len(body) > 900:
            body = body[:900].rstrip() + "..."
        chunks.append(f"### {title}\n{body}")
    return "\n\n".join(chunks)


def sheet_preview(path: Path, sheet_name: str, rows: int = 40) -> pd.DataFrame:
    wb = _load(path, data_only=False)
    if sheet_name not in wb.sheetnames:
        return pd.DataFrame()
    ws = wb[sheet_name]
    data = []
    for row in ws.iter_rows(min_row=1, max_row=min(ws.max_row, rows), max_col=min(ws.max_column, 12), values_only=True):
        cleaned = [_text(value) for value in row]
        if any(cleaned):
            data.append(cleaned)
    width = max((len(row) for row in data), default=0)
    return pd.DataFrame([(row + [""] * width)[:width] for row in data])


def sheet_names(path: Path) -> list[str]:
    wb = _load(path)
    return list(wb.sheetnames)


def export_updated_workbook(log_df: pd.DataFrame | None, source_path: Path) -> str:
    if log_df is None:
        log_df = pd.DataFrame(columns=LOG_COLUMNS)
    log_df = recalculate_log(log_df, source_path)

    with NamedTemporaryFile(prefix="totem_updated_", suffix=".xlsx", delete=False) as handle:
        output_path = Path(handle.name)

    wb = _load(source_path, data_only=False)
    ws = wb["TOTEM_LOG"]

    for row_index in range(7, 87):
        for col_index in list(range(1, 10)) + [14, 15]:
            ws.cell(row_index, col_index).value = None

    for offset, (_, row) in enumerate(log_df.head(80).iterrows(), start=7):
        ws.cell(offset, 1).value = _text(row.get("Sequence"))
        ws.cell(offset, 2).value = _text(row.get("Stanza ID"))
        ws.cell(offset, 3).value = _text(row.get("Draft / Pass"))
        for metric_offset, metric in enumerate(METRICS, start=4):
            ws.cell(offset, metric_offset).value = _number(row.get(metric))
        ws.cell(offset, 14).value = _text(row.get("Priority Fix"))
        ws.cell(offset, 15).value = _text(row.get("Notes"))

    wb.save(output_path)
    return str(output_path)