"""Split source files into retrieval chunks.

Chunks stay close to the natural unit of each format: a documentation section,
one ROS interface definition, or a window of an MJCF file. Every chunk carries
enough header text ("MuJoCo XML reference > geom") that it reads on its own
when pasted into a prompt.
"""

from __future__ import annotations

import hashlib
import re
from dataclasses import dataclass, field
from pathlib import PurePosixPath

MAX_CHARS = 1800
OVERLAP = 200


@dataclass
class Chunk:
    text: str
    source: str  # e.g. "mujoco_docs:doc/XMLreference.rst#geom"
    meta: dict = field(default_factory=dict)

    @property
    def id(self) -> str:
        return hashlib.sha1(f"{self.source}\n{self.text}".encode()).hexdigest()


def _windows(text: str, max_chars: int = MAX_CHARS, overlap: int = OVERLAP) -> list[str]:
    """Split on paragraph boundaries into pieces of at most ~max_chars."""
    text = text.strip()
    if len(text) <= max_chars:
        return [text] if text else []
    paragraphs = re.split(r"\n\s*\n", text)
    out: list[str] = []
    buf = ""
    for para in paragraphs:
        while len(para) > max_chars:  # a single huge paragraph (e.g. a table)
            if buf:
                out.append(buf)
                buf = ""
            out.append(para[:max_chars])
            para = para[max_chars - overlap :]
        if buf and len(buf) + len(para) + 2 > max_chars:
            out.append(buf)
            buf = buf[-overlap:] + "\n\n" + para if overlap else para
        else:
            buf = f"{buf}\n\n{para}" if buf else para
    if buf.strip():
        out.append(buf)
    return out


# RST section titles are underlined (and optionally overlined) with punctuation.
_RST_TITLE = re.compile(r"^(?P<title>[^\n]+)\n(?P<rule>[=\-~^\"'`#*+:.]{3,})[ \t]*$", re.MULTILINE)
_MD_TITLE = re.compile(r"^(#{1,4})\s+(.+)$", re.MULTILINE)


def _sections(text: str, markdown: bool) -> list[tuple[str, str]]:
    """(title, body) pairs in document order; text before the first title gets title ''."""
    pattern = _MD_TITLE if markdown else _RST_TITLE
    matches = list(pattern.finditer(text))
    if not matches:
        return [("", text)]
    sections = []
    if matches[0].start() > 0:
        sections.append(("", text[: matches[0].start()]))
    for i, m in enumerate(matches):
        title = (m.group(2) if markdown else m.group("title")).strip()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
        sections.append((title, text[m.end() : end]))
    return sections


def chunk_document(text: str, source: str, meta: dict, doc_title: str) -> list[Chunk]:
    markdown = source.endswith(".md")
    chunks = []
    for title, body in _sections(text, markdown):
        header = f"{doc_title} > {title}" if title else doc_title
        for piece in _windows(body):
            if len(piece) < 40:
                continue
            anchor = re.sub(r"[^a-z0-9]+", "-", title.lower()).strip("-")
            chunks.append(Chunk(f"[{header}]\n{piece}", f"{source}#{anchor}" if anchor else source, dict(meta)))
    return chunks


def chunk_code(text: str, source: str, meta: dict, label: str) -> list[Chunk]:
    return [Chunk(f"[{label}]\n{piece}", source, dict(meta)) for piece in _windows(text, overlap=0)]


def interface_name(path: str, package: str | None = None) -> str | None:
    """'std_msgs/msg/String.msg' -> 'std_msgs/msg/String' (package can be given for flat repos)."""
    p = PurePosixPath(path)
    if p.suffix not in (".msg", ".srv", ".action") or p.parent.name not in ("msg", "srv", "action"):
        return None
    pkg = package or (p.parent.parent.name if p.parent.parent.name else None)
    if not pkg:
        return None
    return f"{pkg}/{p.parent.name}/{p.stem}"


def interface_kind(name: str) -> str:
    return "Service" if "/srv/" in name else "Action" if "/action/" in name else "Message"


def chunk_interface(text: str, source: str, meta: dict, name: str) -> list[Chunk]:
    kind = interface_kind(name)
    distro = meta.get("distro", "")
    body = text.strip()
    return [Chunk(f"[ROS 2 {distro} {kind} {name}]\n{body}", source, {**meta, "interface": name})]


def chunk_file(path: str, text: str, source_name: str, meta: dict, package: str | None = None) -> list[Chunk]:
    source = f"{source_name}:{path}"
    suffix = PurePosixPath(path).suffix
    if suffix in (".msg", ".srv", ".action"):
        name = interface_name(path, package)
        return chunk_interface(text, source, meta, name) if name else []
    if path.endswith("XMLschema.rst"):
        return [schema_chunk(k, e, source, meta) for k, e in parse_mjcf_schema(text).items()]
    if suffix in (".rst", ".md"):
        return chunk_document(text, source, meta, PurePosixPath(path).stem.replace("-", " "))
    if suffix == ".xml" and "launch" not in path:
        return chunk_code(text, source, meta, f"MJCF example {path}")
    return chunk_code(text, source, meta, f"Example file {path}")


# --------------------------------------------------------------------------- MJCF schema

_DROPDOWN = re.compile(r"\.\. dropdown:: :ref:`([^<`]+)<([\w-]+)>`")
_ATTR_REF = re.compile(r"^\s+:ref:`(\w+)<([\w-]+)>`\s*$")


def parse_mjcf_schema(text: str) -> dict[str, dict]:
    """Parse MuJoCo's doc/XMLschema.rst into {ref_key: {"element", "path", "attributes"}}.

    ref_key is the doc anchor, e.g. "body-geom"; path is the parent chain, e.g. "body/geom".
    """
    schema: dict[str, dict] = {}
    key = None
    for line in text.splitlines():
        m = _DROPDOWN.search(line)
        if m:
            element, key = m.groups()
            element = element.replace("(world)", "")
            schema.setdefault(key, {"element": element, "path": key.replace("-", "/"), "attributes": []})
            continue
        a = _ATTR_REF.match(line)
        if a and key is not None and a.group(2) == f"{key}-{a.group(1)}":
            schema[key]["attributes"].append(a.group(1))
    return schema


def schema_chunk(key: str, entry: dict, source: str, meta: dict) -> Chunk:
    attrs = ", ".join(entry["attributes"]) or "(no attributes)"
    text = f"[MJCF schema] <{entry['element']}> (path {entry['path']}): allowed attributes: {attrs}"
    return Chunk(text, f"{source}#{key}", {**meta, "kind": "schema", "element": entry["element"], "schema_key": key})