"""Sharded long-term memory: hub MEMORY.md + memory/shards/*.md.

Layout modes:
- **Monolith**: no marker file `memory/shards/.mugi_memory_layout`; append_to_memory writes root MEMORY.md only.
- **Sharded**: marker exists; appends go to memory/shards/{slug}.md; hub TOC is regenerated automatically.

Bootstrap: new workspaces get sharded layout when MEMORY.md is first created from template (see persistence.ensure_mugi_workspace).
"""
from __future__ import annotations

import os
import re
from datetime import datetime
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    pass

# Marker file: presence means sharded layout is active.
LAYOUT_MARKER_REL = "memory/shards/.mugi_memory_layout"
DEFAULT_SHARD = "general"
MAX_SHARD_FILES = 200
MAX_ARCHIVE_FILES = 200
MAX_SUMMARY_FILES = 365
MAX_STEM_LEN = 64
MAX_TOC_BLURB_LEN = 120

# Strict slug: lowercase alnum, underscore, hyphen; no path segments.
_SLUG_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{0,63}$")
# memory/archives/MEMORY_ARCHIVE_YYYY-MM.md or MEMORY_ARCHIVE_YYYY-MM__{shard}.md
_ARCHIVE_STEM_RE = re.compile(
    r"^MEMORY_ARCHIVE_\d{4}-\d{2}(?:__[a-z0-9_-]{1,48})?$"
)
_SUMMARY_DAILY_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$")
_SUMMARY_WEEKLY_RE = re.compile(r"^\d{4}-W\d{2}$")

TOC_START = "<!-- mugi:memory-toc:start -->"
TOC_END = "<!-- mugi:memory-toc:end -->"

# Dated sections in MEMORY.md (`## YYYY-MM-DD ...`) usually live in `memory/shards/*.md` after migration.
_DATED_MEMORY_SECTION_RE = re.compile(r"^##\s+\d{4}-\d{2}-\d{2}")


def strip_dated_memory_blocks_from_hub_text(text: str) -> str:
    """Remove `## YYYY-MM-DD` sections from hub MEMORY.md when layout is sharded.

    Keeps preamble and non-dated `##` sections (e.g. TOC, undated notes). Used for
    consolidation preload only, so dated bulk is not duplicated alongside shard reads.
    """
    from utils.memory_text import blocks_to_content, split_into_blocks

    raw = text or ""
    blocks = split_into_blocks(raw)
    if not blocks:
        return raw.strip()
    kept: list[tuple[str, str]] = []
    for header, body in blocks:
        h = (header or "").strip()
        if h and _DATED_MEMORY_SECTION_RE.match(h):
            continue
        kept.append((header, body))
    return blocks_to_content(kept).strip()


def normalize_shard_slug(raw: str | None) -> str | None:
    """Return normalized slug or None if invalid."""
    if raw is None:
        return DEFAULT_SHARD
    s = (raw or "").strip().lower()
    if not s:
        return DEFAULT_SHARD
    # Strip accidental .md
    if s.endswith(".md"):
        s = s[:-3]
    if not _SLUG_RE.match(s):
        return None
    if s in (".", "..") or "/" in s or "\\" in s:
        return None
    return s


def default_shard() -> str:
    return DEFAULT_SHARD


def layout_marker_abs(workspace: str) -> str:
    return os.path.join(workspace, LAYOUT_MARKER_REL)


def layout_is_sharded(workspace: str) -> bool:
    return os.path.isfile(layout_marker_abs(workspace))


def shard_rel_path(slug: str) -> str:
    """Relative path memory/shards/{slug}.md (slug must be pre-validated)."""
    return f"memory/shards/{slug}.md"


def list_shard_rel_paths(workspace: str) -> list[str]:
    """Workspace-relative paths memory/shards/*.md, sorted, bounded; skips dotfiles and symlinks."""
    sd = os.path.join(workspace, "memory", "shards")
    if not os.path.isdir(sd):
        return []
    out: list[str] = []
    try:
        names = sorted(
            f for f in os.listdir(sd)
            if f.endswith(".md") and not f.startswith(".") and len(f) <= MAX_STEM_LEN + 3
        )
    except OSError:
        return []
    for name in names[:MAX_SHARD_FILES]:
        abs_p = os.path.join(sd, name)
        if os.path.isfile(abs_p) and not os.path.islink(abs_p):
            out.append(f"memory/shards/{name}")
    return out


def list_shard_abs_paths(workspace: str) -> list[str]:
    return [os.path.join(workspace, r) for r in list_shard_rel_paths(workspace)]


def validate_archive_stem(raw: str | None) -> str | None:
    """Return archive filename stem (no .md) or None if invalid."""
    if raw is None:
        return None
    s = (raw or "").strip()
    if s.endswith(".md"):
        s = s[:-3]
    if not s or "/" in s or "\\" in s or s in (".", ".."):
        return None
    if not _ARCHIVE_STEM_RE.match(s):
        return None
    return s


def validate_summary_slug(raw: str | None) -> str | None:
    """Return summary filename stem (daily or weekly) or None if invalid."""
    if raw is None:
        return None
    s = (raw or "").strip()
    if s.endswith(".md"):
        s = s[:-3]
    if not s or "/" in s or "\\" in s or s in (".", ".."):
        return None
    if _SUMMARY_DAILY_RE.match(s):
        try:
            datetime.strptime(s, "%Y-%m-%d")
        except ValueError:
            return None
        return s
    if _SUMMARY_WEEKLY_RE.match(s):
        return s
    return None


def archive_rel_path(stem: str) -> str:
    """Relative path memory/archives/{stem}.md (stem must be pre-validated)."""
    return f"memory/archives/{stem}.md"


def summary_rel_path(slug: str) -> str:
    """Relative path memory/summaries/{slug}.md (slug must be pre-validated)."""
    return f"memory/summaries/{slug}.md"


def list_archive_rel_paths(workspace: str) -> list[str]:
    """Workspace-relative paths memory/archives/MEMORY_ARCHIVE_*.md, newest first."""
    ad = os.path.join(workspace, "memory", "archives")
    if not os.path.isdir(ad):
        return []
    out: list[str] = []
    try:
        names = sorted(
            (
                f
                for f in os.listdir(ad)
                if f.startswith("MEMORY_ARCHIVE_")
                and f.endswith(".md")
                and not f.startswith(".")
            ),
            reverse=True,
        )
    except OSError:
        return []
    for name in names[:MAX_ARCHIVE_FILES]:
        stem = name[:-3]
        if not validate_archive_stem(stem):
            continue
        abs_p = os.path.join(ad, name)
        if os.path.isfile(abs_p) and not os.path.islink(abs_p):
            out.append(f"memory/archives/{name}")
    return out


def list_archive_stems(workspace: str) -> list[str]:
    """Archive filename stems (no .md), sorted newest first."""
    return [os.path.basename(r)[:-3] for r in list_archive_rel_paths(workspace)]


def list_summary_rel_paths(workspace: str) -> list[str]:
    """Workspace-relative paths memory/summaries/*.md (daily + weekly), newest first."""
    sd = os.path.join(workspace, "memory", "summaries")
    if not os.path.isdir(sd):
        return []
    out: list[str] = []
    try:
        names = sorted(
            (
                f
                for f in os.listdir(sd)
                if f.endswith(".md")
                and not f.startswith(".")
                and (
                    _SUMMARY_DAILY_RE.match(f[:-3])
                    or _SUMMARY_WEEKLY_RE.match(f[:-3])
                )
            ),
            reverse=True,
        )
    except OSError:
        return []
    for name in names[:MAX_SUMMARY_FILES]:
        slug = name[:-3]
        if not validate_summary_slug(slug):
            continue
        abs_p = os.path.join(sd, name)
        if os.path.isfile(abs_p) and not os.path.islink(abs_p):
            out.append(f"memory/summaries/{name}")
    return out


def list_summary_slugs(workspace: str) -> list[str]:
    """Summary filename stems (no .md), sorted newest first."""
    return [os.path.basename(r)[:-3] for r in list_summary_rel_paths(workspace)]


def _shard_title_and_blurb(abs_path: str) -> tuple[str, str]:
    from utils.memory_text import read_memory_file

    text = read_memory_file(abs_path)
    base = os.path.basename(abs_path)
    stem = base[:-3] if base.endswith(".md") else base
    title = stem.replace("_", " ").replace("-", " ").title()
    blurb = ""
    # Skip YAML frontmatter rough scan
    body = text
    if body.strip().startswith("---"):
        parts = body.split("---", 2)
        if len(parts) >= 3:
            body = parts[2]
    for line in body.splitlines():
        ln = line.strip()
        if not ln or ln.startswith("#"):
            continue
        blurb = ln[:MAX_TOC_BLURB_LEN]
        if len(ln) > MAX_TOC_BLURB_LEN:
            blurb = blurb.rstrip() + "…"
        break
    return title, blurb


def regenerate_memory_hub_toc(workspace: str) -> str | None:
    """Rewrite only the TOC region in MEMORY.md. Returns error string or None on success."""
    from utils.path_safety import safe_resolve_path, secure_write_file
    from utils.memory_text import read_memory_file

    mem_resolved, err = safe_resolve_path(workspace, "MEMORY.md", allow_write=True, create_parents=True)
    if err or not mem_resolved:
        return err or "Could not resolve MEMORY.md"
    content = read_memory_file(mem_resolved)
    rels = list_shard_rel_paths(workspace)
    toc_lines = ["## Memory shards", ""]
    for rel in rels:
        abs_p = os.path.join(workspace, rel)
        title, blurb = _shard_title_and_blurb(abs_p)
        link = rel.replace(" ", "%20")
        if blurb:
            toc_lines.append(f"- [{title}]({link}) — {blurb}")
        else:
            toc_lines.append(f"- [{title}]({link})")
    toc_lines.append("")
    toc_body = "\n".join(toc_lines)
    new_toc_section = f"{TOC_START}\n{toc_body}{TOC_END}"

    if TOC_START in content and TOC_END in content:
        pre, rest = content.split(TOC_START, 1)
        _, post = rest.split(TOC_END, 1)
        new_content = pre.rstrip() + "\n\n" + new_toc_section + "\n" + post.lstrip()
    else:
        new_content = (content.rstrip() + "\n\n" + new_toc_section + "\n").strip() + "\n"

    try:
        secure_write_file(mem_resolved, new_content, mode=0o600)
    except Exception as e:
        return str(e)
    return None


def write_layout_marker(workspace: str) -> None:
    os.makedirs(os.path.join(workspace, "memory", "shards"), exist_ok=True)
    path = layout_marker_abs(workspace)
    if not os.path.isfile(path):
        with open(path, "w", encoding="utf-8") as f:
            f.write("v1\n")
        os.chmod(path, 0o600)


def ensure_shard_file(workspace: str, slug: str) -> tuple[str | None, str]:
    """Create memory/shards/{slug}.md with minimal frontmatter if missing."""
    from utils.path_safety import safe_resolve_path, secure_write_file
    from utils.memory_text import read_memory_file

    normalized = normalize_shard_slug(slug)
    if normalized is None:
        return None, "Invalid shard slug: use letters, digits, underscore, hyphen (max 64 chars)."
    rel = shard_rel_path(normalized)
    abs_path, err = safe_resolve_path(workspace, rel, allow_write=True, create_parents=True)
    if err or not abs_path:
        return None, err or "resolve failed"
    if os.path.isfile(abs_path) and read_memory_file(abs_path).strip():
        return abs_path, ""
    title = normalized.replace("_", " ").replace("-", " ").title()
    seed = (
        "---\n"
        f"id: mem_{normalized}\n"
        f"title: {title}\n"
        "---\n\n"
        f"# {title}\n\n"
    )
    try:
        secure_write_file(abs_path, seed, mode=0o600)
    except Exception as e:
        return None, str(e)
    return abs_path, ""


def ensure_general_shard(workspace: str) -> tuple[str | None, str]:
    """Ensure memory/shards/general.md exists (minimal). Returns (abs_path, err)."""
    from utils.path_safety import safe_resolve_path, secure_write_file
    from utils.memory_text import read_memory_file

    rel = shard_rel_path(DEFAULT_SHARD)
    abs_path, err = safe_resolve_path(workspace, rel, allow_write=True, create_parents=True)
    if err or not abs_path:
        return None, err or "resolve failed"
    if not os.path.isfile(abs_path) or not read_memory_file(abs_path).strip():
        seed = (
            "---\n"
            f"id: mem_{DEFAULT_SHARD}\n"
            f"title: General\n"
            "---\n\n"
            f"# General memory\n\n"
            f"Long-term facts default to this shard. Use **append_to_memory** with optional `shard` "
            f"for other topics.\n"
        )
        try:
            secure_write_file(abs_path, seed, mode=0o600)
        except Exception as e:
            return None, str(e)
    return abs_path, ""


def ensure_sharded_bootstrap_new_workspace(workspace_path: str) -> None:
    """Called when MEMORY.md was just created from template: activate sharded layout."""
    if layout_is_sharded(workspace_path):
        ensure_general_shard(workspace_path)
        regenerate_memory_hub_toc(workspace_path)
        return
    # Do not bootstrap if user already has a fat monolith MEMORY (e.g. copied workspace)
    mem_path = os.path.join(workspace_path, "MEMORY.md")
    if os.path.isfile(mem_path):
        try:
            if os.path.getsize(mem_path) > 8000:
                return
        except OSError:
            pass
    write_layout_marker(workspace_path)
    ensure_general_shard(workspace_path)
    regenerate_memory_hub_toc(workspace_path)


def migrate_memory_to_shards_impl(workspace: str, *, force: bool = False) -> str:
    """Move dated blocks from MEMORY.md to general.md; write hub + marker. Idempotent."""
    from utils.path_safety import safe_resolve_path, secure_write_file
    from utils.memory_text import read_memory_file, split_into_blocks, blocks_to_content

    if layout_is_sharded(workspace) and not force:
        return "Already using sharded memory layout (marker present). Nothing to migrate."

    mem_path, err = safe_resolve_path(workspace, "MEMORY.md", allow_write=False, create_parents=False)
    if err or not mem_path or not os.path.isfile(mem_path):
        return "MEMORY.md not found; nothing to migrate."
    content = read_memory_file(mem_path)
    blocks = split_into_blocks(content)

    # Dated memory blocks: ## YYYY-MM-DD ...
    dated: list[tuple[str, str]] = []
    hub_parts: list[tuple[str, str]] = []
    for header, body in blocks:
        h = header or ""
        if re.search(r"##\s+\d{4}-\d{2}-\d{2}", h):
            dated.append((header, body))
        else:
            hub_parts.append((header, body))

    if not dated and content.strip():
        dated = [
            (
                f"## {datetime.now().strftime('%Y-%m-%d %H:%M')} (imported monolith)",
                content.strip(),
            )
        ]
        hub_parts = []

    gen_path, gerr = ensure_general_shard(workspace)
    if gerr or not gen_path:
        return f"Migration failed: could not ensure general shard: {gerr}"

    existing_general = read_memory_file(gen_path)
    general_blocks = split_into_blocks(existing_general) if existing_general.strip() else []
    # Merge: seed frontmatter blocks may be in general; append dated from MEMORY
    merged = general_blocks + dated
    new_general = blocks_to_content(merged) if merged else existing_general
    try:
        secure_write_file(gen_path, new_general.strip() + "\n" if new_general.strip() else existing_general, mode=0o600)
    except Exception as e:
        return f"Migration failed writing general.md: {e}"

    # Hub: non-dated content + TOC
    hub_core = blocks_to_content(hub_parts).strip()
    if not hub_core:
        hub_core = (
            "# MEMORY.md — Long-term memory hub\n\n"
            "Curated index of memory shards. Details live under `memory/shards/`.\n"
        )
    try:
        secure_write_file(mem_path, hub_core.rstrip() + "\n", mode=0o600)
    except Exception as e:
        return f"Migration failed trimming MEMORY.md: {e}"
    toc_err = regenerate_memory_hub_toc(workspace)
    if toc_err:
        return f"Migrated but hub TOC failed: {toc_err}"
    write_layout_marker(workspace)
    n_dated = len(dated)
    return (
        f"Migrated {n_dated} block(s) to memory/shards/general.md; "
        "MEMORY.md is now a hub. Sharded layout is active."
    )
