from __future__ import annotations

import json
import re
import zipfile
from pathlib import Path
from xml.etree import ElementTree

import pymupdf4llm

RUBRIC_EXTENSIONS = {".pdf", ".docx", ".txt", ".md"}
STUDENT_EXTENSIONS = RUBRIC_EXTENSIONS | {".ipynb"}


def supported_extensions(for_student: bool = False) -> set[str]:
    return STUDENT_EXTENSIONS if for_student else RUBRIC_EXTENSIONS


def is_supported(filename: str, *, for_student: bool = False) -> bool:
    return Path(filename).suffix.lower() in supported_extensions(for_student)


def supported_extensions_label(*, for_student: bool = False) -> str:
    return ", ".join(sorted(supported_extensions(for_student)))


def clean_text(text: str) -> str:
    text = text.replace("\x00", " ")
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def _read_plain_text(path: Path) -> str:
    return path.read_text("utf-8", errors="replace")


def _xml_text(xml_bytes: bytes) -> str:
    root = ElementTree.fromstring(xml_bytes)
    ns = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}
    blocks: list[str] = []
    for paragraph in root.findall(".//w:p", ns):
        text = "".join(node.text or "" for node in paragraph.findall(".//w:t", ns))
        if text.strip():
            blocks.append(text)
    return "\n".join(blocks)


def _read_docx(path: Path) -> str:
    parts = [
        "word/document.xml",
        "word/comments.xml",
        "word/footnotes.xml",
        "word/endnotes.xml",
    ]
    with zipfile.ZipFile(path) as archive:
        text_parts = []
        for part in parts:
            try:
                text = _xml_text(archive.read(part))
            except KeyError:
                continue
            if text.strip():
                text_parts.append(text)
        return "\n\n".join(text_parts)


def _read_ipynb(path: Path) -> str:
    try:
        import nbformat
        from nbconvert import MarkdownExporter

        notebook = nbformat.read(path, as_version=4)
        exporter = MarkdownExporter()
        body, _ = exporter.from_notebook_node(notebook)
        return body
    except Exception:
        data = json.loads(path.read_text("utf-8", errors="replace"))
        blocks = []
        for cell in data.get("cells", []):
            source = cell.get("source", "")
            if isinstance(source, list):
                source = "".join(source)
            if source:
                blocks.append(source)
        return "\n\n".join(blocks)


def parse_file_text(path: Path, filename: str | None = None) -> str:
    name = filename or path.name
    suffix = Path(name).suffix.lower()
    if suffix == ".pdf":
        text = pymupdf4llm.to_markdown(str(path))
    elif suffix == ".docx":
        text = _read_docx(path)
    elif suffix in {".txt", ".md"}:
        text = _read_plain_text(path)
    elif suffix == ".ipynb":
        text = _read_ipynb(path)
    else:
        raise ValueError(f"Unsupported file type: {suffix or name}")
    return clean_text(text)