#!/usr/bin/env python3
"""Prepare Potato annotation data from the math annotator Excel workbook.

This script intentionally uses only the Python standard library so it can run
in a fresh environment without openpyxl/pandas.
"""

from __future__ import annotations

import argparse
import csv
import html
import json
import re
import sys
import textwrap
import zipfile
from collections import defaultdict
from pathlib import Path
from xml.etree import ElementTree as ET


ROOT = Path(__file__).resolve().parents[2]
DEFAULT_INPUT = ROOT / "annotated_data" / "math_annotator_train&test_sets_simple_2q_train_2q_test.xlsx"
DEFAULT_OUTPUT_DIR = Path(__file__).resolve().parents[1] / "my-annotation-task" / "data"

SHEET_TO_SPLIT = {
    "math_annotator_training_set": "train",
    "math_annotator_testing_set": "test",
}

DISPLAY_COLUMNS = ["dialog_context", "correct_solution", "tutor_response"]

DIMENSIONS = [
    "Content Correctness",
    "Learner-State Assessment",
    "Issue Localization",
    "Disclosure Appropriateness",
    "Providing Guidance",
    "Coherence",
    "Actionability",
    "Clarity",
    "Conciseness",
    "Humanness",
]

VALID_LABELS = {"Yes", "To some extent", "No"}

NS = {
    "a": "http://schemas.openxmlformats.org/spreadsheetml/2006/main",
    "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships",
}


def column_index(cell_ref: str) -> int:
    match = re.match(r"([A-Z]+)", cell_ref or "A")
    if not match:
        return 0
    index = 0
    for char in match.group(1):
        index = index * 26 + (ord(char) - ord("A") + 1)
    return index - 1


def load_shared_strings(archive: zipfile.ZipFile) -> list[str]:
    if "xl/sharedStrings.xml" not in archive.namelist():
        return []
    root = ET.fromstring(archive.read("xl/sharedStrings.xml"))
    strings: list[str] = []
    for string_item in root.findall("a:si", NS):
        strings.append("".join(node.text or "" for node in string_item.findall(".//a:t", NS)))
    return strings


def get_cell_text(cell: ET.Element, shared_strings: list[str]) -> str:
    cell_type = cell.attrib.get("t")
    value_node = cell.find("a:v", NS)
    if cell_type == "s" and value_node is not None and value_node.text:
        return shared_strings[int(value_node.text)]
    if cell_type == "inlineStr":
        return "".join(node.text or "" for node in cell.findall(".//a:t", NS))
    if value_node is not None:
        return value_node.text or ""
    return ""


def resolve_sheet_path(target: str) -> str:
    target = target.lstrip("/")
    if target.startswith("xl/"):
        return target
    return f"xl/{target}"


def read_workbook(path: Path) -> dict[str, list[list[str]]]:
    with zipfile.ZipFile(path) as archive:
        shared_strings = load_shared_strings(archive)
        workbook_root = ET.fromstring(archive.read("xl/workbook.xml"))
        rels_root = ET.fromstring(archive.read("xl/_rels/workbook.xml.rels"))
        relationship_targets = {
            rel.attrib["Id"]: rel.attrib["Target"]
            for rel in rels_root
        }

        sheets: dict[str, list[list[str]]] = {}
        for sheet in workbook_root.find("a:sheets", NS):
            sheet_name = sheet.attrib["name"]
            rel_id = sheet.attrib[f"{{{NS['r']}}}id"]
            sheet_path = resolve_sheet_path(relationship_targets[rel_id])
            sheet_root = ET.fromstring(archive.read(sheet_path))
            rows: list[list[str]] = []

            for row in sheet_root.findall(".//a:sheetData/a:row", NS):
                values_by_col = {
                    column_index(cell.attrib.get("r", "A")): get_cell_text(cell, shared_strings)
                    for cell in row.findall("a:c", NS)
                }
                if values_by_col:
                    width = max(values_by_col) + 1
                    rows.append([values_by_col.get(col, "") for col in range(width)])
                else:
                    rows.append([])
            sheets[sheet_name] = rows
    return sheets


def normalize_rows(raw_rows: list[list[str]]) -> list[dict[str, str]]:
    nonempty_rows = [row for row in raw_rows if any(str(value).strip() for value in row)]
    if not nonempty_rows:
        return []

    header = [str(value).strip() for value in nonempty_rows[0]]
    rows: list[dict[str, str]] = []
    for raw_row in nonempty_rows[1:]:
        padded = raw_row + [""] * (len(header) - len(raw_row))
        rows.append({header[index]: str(padded[index]).strip() for index in range(len(header))})
    return rows


def slugify(value: str) -> str:
    value = value.lower().strip()
    value = re.sub(r"[^a-z0-9]+", "_", value)
    return value.strip("_") or "item"


def escape_text_node(text: str) -> str:
    """Escape HTML text content without turning quotes into visible entities."""
    return html.escape(html.unescape(text), quote=False)


def inline_format(text: str) -> str:
    escaped = escape_text_node(text)
    escaped = re.sub(r"`([^`]+)`", r"<code>\1</code>", escaped)
    escaped = re.sub(r"\*\*([^*]+)\*\*", r"<strong>\1</strong>", escaped)
    escaped = re.sub(
        r"(?m)^(Tutor|Student):",
        lambda match: f"<strong>{match.group(1)}:</strong>",
        escaped,
    )
    return escaped


def rich_text(text: str) -> str:
    text = (text or "").strip()
    if not text:
        return "<p><em>No content provided.</em></p>"

    parts = re.split(r"(```(?:[a-zA-Z0-9_+-]+)?\n.*?\n```)", text, flags=re.DOTALL)
    rendered: list[str] = []
    for part in parts:
        if not part:
            continue
        fence_match = re.match(r"```(?:[a-zA-Z0-9_+-]+)?\n(.*?)\n```", part, flags=re.DOTALL)
        if fence_match:
            code = escape_text_node(fence_match.group(1).strip("\n"))
            rendered.append(f'<pre class="potato-code-block"><code>{code}</code></pre>')
            continue

        paragraphs = [paragraph.strip() for paragraph in re.split(r"\n\s*\n", part) if paragraph.strip()]
        for paragraph in paragraphs:
            paragraph_html = inline_format(paragraph).replace("\n", "<br>")
            rendered.append(f"<p>{paragraph_html}</p>")
    return "\n".join(rendered)


def section_html(title: str, body: str, *, extra_class: str = "") -> str:
    class_name = "potato-text-section"
    if extra_class:
        class_name = f"{class_name} {extra_class}"
    return (
        f'<div class="{class_name}">'
        f"<h3>{html.escape(title)}</h3>"
        f"{rich_text(body)}"
        "</div>"
    )


def make_text2show(row: dict[str, str]) -> str:
    return textwrap.dedent(
        f"""\
        Dialog context:
        {row["dialog_context"]}

        Reference (correct) solution:
        {row["correct_solution"]}

        Next tutor response:
        {row["tutor_response"]}
        """
    ).strip()


def make_text2show_html(row: dict[str, str], split: str, item_id: str) -> str:
    split_label = split.upper()
    return "\n".join(
        [
            f'<div class="potato-instance-meta"><span>{split_label}</span><span>{html.escape(item_id)}</span></div>',
            section_html("Dialog context", row["dialog_context"]),
            section_html("Reference (correct) solution", row["correct_solution"], extra_class="reference-solution"),
            section_html("Next tutor response", row["tutor_response"], extra_class="tutor-response"),
        ]
    )


def validate_labels(row: dict[str, str], row_id: str) -> None:
    for dimension in DIMENSIONS:
        label = row.get(dimension, "").strip()
        if label not in VALID_LABELS:
            raise ValueError(f"{row_id}: invalid label for {dimension!r}: {label!r}")


def prepare_records(workbook_path: Path) -> dict[str, list[dict[str, str]]]:
    sheets = read_workbook(workbook_path)
    missing_sheets = [sheet for sheet in SHEET_TO_SPLIT if sheet not in sheets]
    if missing_sheets:
        raise ValueError(f"Missing expected sheets: {missing_sheets}")

    records_by_split: dict[str, list[dict[str, str]]] = {}
    for sheet_name, split in SHEET_TO_SPLIT.items():
        rows = normalize_rows(sheets[sheet_name])
        if not rows:
            raise ValueError(f"Sheet {sheet_name!r} has no data rows")

        required = DISPLAY_COLUMNS + DIMENSIONS
        missing_columns = sorted({column for column in required if column not in rows[0]})
        if missing_columns:
            raise ValueError(f"Sheet {sheet_name!r} missing columns: {missing_columns}")

        dialog_numbers: dict[str, int] = {}
        response_counts: defaultdict[int, int] = defaultdict(int)
        split_records: list[dict[str, str]] = []

        for row in rows:
            dialog_context = row["dialog_context"]
            if dialog_context not in dialog_numbers:
                dialog_numbers[dialog_context] = len(dialog_numbers) + 1
            dialog_id = dialog_numbers[dialog_context]
            response_counts[dialog_id] += 1
            response_id = response_counts[dialog_id]
            item_id = f"math_{split}_q{dialog_id:02d}_r{response_id:02d}"

            validate_labels(row, item_id)

            record = {
                "id": item_id,
                "split": split,
                "domain": "math",
                "dialog_id": f"q{dialog_id:02d}",
                "response_id": f"r{response_id:02d}",
                "text2show": make_text2show(row),
                "text2show_html": make_text2show_html(row, split, item_id),
            }
            for column in DISPLAY_COLUMNS + DIMENSIONS:
                record[column] = row[column]
            split_records.append(record)
        records_by_split[split] = split_records
    return records_by_split


def write_csv(path: Path, rows: list[dict[str, str]]) -> None:
    fieldnames = [
        "id",
        "split",
        "domain",
        "dialog_id",
        "response_id",
        "text2show",
        "text2show_html",
        *DISPLAY_COLUMNS,
        *DIMENSIONS,
    ]
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def to_gold_item(record: dict[str, str], *, key_name: str) -> dict[str, object]:
    labels = {dimension: record[dimension] for dimension in DIMENSIONS}
    return {
        "id": record["id"],
        "text": record["text2show"],
        "text2show": record["text2show"],
        "text2show_html": record["text2show_html"],
        key_name: labels,
        "explanation": "Gold labels are provided by the curated math annotator training/test workbook.",
        "metadata": {
            "split": record["split"],
            "domain": record["domain"],
            "dialog_id": record["dialog_id"],
            "response_id": record["response_id"],
        },
    }


def write_json(path: Path, payload: object) -> None:
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, default=DEFAULT_INPUT, help="Source .xlsx workbook")
    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR, help="Output data directory")
    args = parser.parse_args()

    if not args.input.exists():
        print(f"Input workbook not found: {args.input}", file=sys.stderr)
        return 1

    records_by_split = prepare_records(args.input)
    output_dir = args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)

    all_records = records_by_split["train"] + records_by_split["test"]
    write_csv(output_dir / "math_annotator_training_set_with_id_text2show.csv", records_by_split["train"])
    write_csv(output_dir / "math_annotator_testing_set_with_id_text2show.csv", records_by_split["test"])
    write_csv(output_dir / "math_annotator_demo_all_with_id_text2show.csv", all_records)

    write_json(
        output_dir / "training_questions.json",
        [to_gold_item(record, key_name="correct_answers") for record in records_by_split["train"]],
    )
    write_json(
        output_dir / "gold_standards.json",
        [to_gold_item(record, key_name="gold_label") for record in records_by_split["test"]],
    )

    summary = {
        "input": str(args.input),
        "outputs": {
            "train_csv": "math_annotator_training_set_with_id_text2show.csv",
            "test_csv": "math_annotator_testing_set_with_id_text2show.csv",
            "combined_csv": "math_annotator_demo_all_with_id_text2show.csv",
            "training_questions": "training_questions.json",
            "gold_standards": "gold_standards.json",
        },
        "counts": {split: len(records) for split, records in records_by_split.items()},
        "dimensions": DIMENSIONS,
        "labels": sorted(VALID_LABELS),
    }
    write_json(output_dir / "data_summary.json", summary)

    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())