"""
ShelfAI Memory Compactor — Consolidates memory files to prevent bloat.

Two modes:
  1. Heuristic (free) — exact/near-duplicate removal, date-based archival,
     entry counting, token estimation. Safe to run anytime.
  2. LLM-assisted (~$0.01/file) — semantic dedup, consolidation of related
     entries, summarization of old patterns. Requires a provider.

Memory files grow over time as the session agent appends lessons, patterns,
and preferences. Without compaction, files accumulate duplicates, superseded
entries, and stale observations that dilute context quality.
"""

import re
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional


__all__ = [
    "scan_memory",
    "compact_heuristic",
    "MemoryStats",
    "CompactResult",
    "MemoryEntry",
]


# ──────────────────────────────────────────────
# Data Classes
# ──────────────────────────────────────────────


@dataclass
class MemoryEntry:
    """A single entry (bullet point) extracted from a memory file."""

    text: str
    section: str  # H2 or H3 heading it lives under
    line_number: int
    date_tag: Optional[str] = None  # extracted date if present
    confidence: Optional[str] = None  # high/medium/low if tagged
    pattern_ref: Optional[str] = None  # e.g. "Pattern H-A"
    is_placeholder: bool = False  # True if "_(pending...)" or "_(none...)"


@dataclass
class MemoryStats:
    """Statistics for a single memory file."""

    path: Path
    total_entries: int
    placeholder_entries: int
    dated_entries: int
    stale_entries: int  # entries older than stale_days
    duplicate_groups: int  # groups of near-duplicate entries
    estimated_tokens: int
    line_count: int
    sections: list  # list of {name, entry_count}
    needs_compaction: bool  # True if file exceeds thresholds


@dataclass
class ScanResult:
    """Results from scanning all memory files."""

    files: list  # list of MemoryStats
    total_entries: int = 0
    total_tokens: int = 0
    files_needing_compaction: int = 0

    def __post_init__(self):
        self.total_entries = sum(f.total_entries for f in self.files)
        self.total_tokens = sum(f.estimated_tokens for f in self.files)
        self.files_needing_compaction = sum(1 for f in self.files if f.needs_compaction)


@dataclass
class CompactResult:
    """Result of compacting a memory file."""

    path: Path
    original_entries: int
    compacted_entries: int
    removed_duplicates: int
    removed_placeholders: int
    archived_stale: int
    original_lines: int
    compacted_lines: int
    original_tokens: int
    compacted_tokens: int
    content: str  # the compacted file content
    changes: list = field(default_factory=list)  # human-readable change descriptions

    @property
    def savings_pct(self) -> int:
        if self.original_tokens == 0:
            return 0
        return round((1 - self.compacted_tokens / self.original_tokens) * 100)


# ──────────────────────────────────────────────
# Token Estimation
# ──────────────────────────────────────────────


def _estimate_tokens(text: str) -> int:
    """Rough token estimate (~4 chars per token for English markdown)."""
    return max(1, len(text) // 4)


# ──────────────────────────────────────────────
# Entry Parsing
# ──────────────────────────────────────────────


_DATE_PATTERN = re.compile(r"(\d{4}-\d{2}-\d{2})")
_CONFIDENCE_PATTERN = re.compile(r"\[confidence:\s*(high|medium|low)\]", re.IGNORECASE)
_PATTERN_REF = re.compile(r"\(Pattern\s+[A-Z]-[A-Z0-9]+[^)]*\)", re.IGNORECASE)
_PLACEHOLDER_PATTERN = re.compile(
    r"^\s*_\s*\((none|pending|pre-launch|no\s+\w+\s+data)[^)]*\)\s*_?\s*$",
    re.IGNORECASE,
)


def _parse_entries(content: str) -> list:
    """Parse a memory file into structured entries."""
    entries = []
    current_section = "top"
    lines = content.split("\n")

    for i, line in enumerate(lines, 1):
        stripped = line.strip()

        # Track sections
        if stripped.startswith("## "):
            current_section = stripped.lstrip("#").strip()
            continue
        if stripped.startswith("### "):
            current_section = stripped.lstrip("#").strip()
            continue

        # Skip non-entry lines
        if not stripped.startswith("- ") and not stripped.startswith("* "):
            continue

        entry_text = stripped[2:].strip()
        if not entry_text:
            continue

        # Extract metadata
        date_match = _DATE_PATTERN.search(entry_text)
        conf_match = _CONFIDENCE_PATTERN.search(entry_text)
        pat_match = _PATTERN_REF.search(entry_text)
        is_placeholder = bool(_PLACEHOLDER_PATTERN.match(entry_text))

        entries.append(MemoryEntry(
            text=entry_text,
            section=current_section,
            line_number=i,
            date_tag=date_match.group(1) if date_match else None,
            confidence=conf_match.group(1).lower() if conf_match else None,
            pattern_ref=pat_match.group(0) if pat_match else None,
            is_placeholder=is_placeholder,
        ))

    return entries


def _normalize_for_dedup(text: str) -> str:
    """Normalize entry text for near-duplicate detection."""
    text = text.lower()
    # Remove dates, confidence tags, pattern refs
    text = _DATE_PATTERN.sub("", text)
    text = _CONFIDENCE_PATTERN.sub("", text)
    text = _PATTERN_REF.sub("", text)
    # Remove punctuation and extra whitespace
    text = re.sub(r"[^a-z0-9\s]", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def _find_duplicate_groups(entries: list) -> list:
    """Find groups of near-duplicate entries. Returns list of (keeper_idx, duplicate_idxs)."""
    normalized = [(i, _normalize_for_dedup(e.text)) for i, e in enumerate(entries)
                  if not e.is_placeholder]

    groups = []
    seen = set()

    for i, (idx_a, norm_a) in enumerate(normalized):
        if idx_a in seen:
            continue
        duplicates = []
        for j, (idx_b, norm_b) in enumerate(normalized[i + 1:], i + 1):
            if idx_b in seen:
                continue
            # Exact normalized match or one is a substring of the other
            if norm_a == norm_b or (len(norm_a) > 20 and norm_a in norm_b) or (len(norm_b) > 20 and norm_b in norm_a):
                duplicates.append(idx_b)
                seen.add(idx_b)
        if duplicates:
            groups.append((idx_a, duplicates))
            seen.add(idx_a)

    return groups


# ──────────────────────────────────────────────
# Scan
# ──────────────────────────────────────────────


def scan_memory(
    shelf_path: Path,
    agents_dir: Optional[Path] = None,
    stale_days: int = 90,
    entry_threshold: int = 30,
    token_threshold: int = 2000,
) -> ScanResult:
    """
    Scan all memory files for compaction candidates.

    Checks both shelf/memory/ files and (optionally) agents/*/MEMORY.md files.
    """
    files = []
    today = datetime.now()
    stale_cutoff = today - timedelta(days=stale_days)

    # Collect memory file paths
    memory_paths = []

    # Shelf memory files
    shelf_memory = shelf_path / "memory"
    if shelf_memory.exists():
        for md in shelf_memory.rglob("*.md"):
            if md.name == ".gitkeep":
                continue
            memory_paths.append(md)

    # Agent MEMORY.md files
    if agents_dir and agents_dir.exists():
        for agent_dir in sorted(agents_dir.iterdir()):
            if not agent_dir.is_dir():
                continue
            mem_file = agent_dir / "MEMORY.md"
            if mem_file.exists():
                memory_paths.append(mem_file)

    for path in sorted(memory_paths):
        content = path.read_text(encoding="utf-8")
        entries = _parse_entries(content)
        line_count = content.count("\n") + 1
        tokens = _estimate_tokens(content)

        # Count stale entries
        stale_count = 0
        for e in entries:
            if e.date_tag:
                try:
                    entry_date = datetime.strptime(e.date_tag, "%Y-%m-%d")
                    if entry_date < stale_cutoff:
                        stale_count += 1
                except ValueError:
                    pass

        # Find duplicates
        dup_groups = _find_duplicate_groups(entries)

        # Build section summary
        section_counts = {}
        for e in entries:
            section_counts[e.section] = section_counts.get(e.section, 0) + 1
        sections = [{"name": k, "entry_count": v} for k, v in section_counts.items()]

        real_entries = sum(1 for e in entries if not e.is_placeholder)
        placeholder_count = sum(1 for e in entries if e.is_placeholder)
        dated_count = sum(1 for e in entries if e.date_tag)

        needs_compaction = (
            real_entries > entry_threshold
            or tokens > token_threshold
            or len(dup_groups) > 0
            or stale_count > 5
            or placeholder_count > 0
        )

        files.append(MemoryStats(
            path=path,
            total_entries=real_entries,
            placeholder_entries=placeholder_count,
            dated_entries=dated_count,
            stale_entries=stale_count,
            duplicate_groups=len(dup_groups),
            estimated_tokens=tokens,
            line_count=line_count,
            sections=sections,
            needs_compaction=needs_compaction,
        ))

    return ScanResult(files=files)


# ──────────────────────────────────────────────
# Heuristic Compaction
# ──────────────────────────────────────────────


def compact_heuristic(
    filepath: Path,
    stale_days: int = 90,
    remove_placeholders: bool = True,
    remove_duplicates: bool = True,
    archive_stale: bool = True,
) -> CompactResult:
    """
    Run heuristic compaction on a memory file. Free, no LLM needed.

    What it does:
    - Removes placeholder entries ("_(pending...)", "_(none...)")
    - Removes near-duplicate entries (keeps the longer/more recent one)
    - Moves entries older than stale_days to a "## Archived" section
    - Preserves file structure (headings, tables, frontmatter)

    Returns a CompactResult with the new content. Does NOT write the file —
    caller decides whether to apply.
    """
    content = filepath.read_text(encoding="utf-8")
    entries = _parse_entries(content)
    original_lines = content.count("\n") + 1
    original_tokens = _estimate_tokens(content)

    changes = []
    remove_lines = set()  # line numbers to remove
    archive_entries = []  # entries to move to Archived section

    today = datetime.now()
    stale_cutoff = today - timedelta(days=stale_days)

    # 1. Remove placeholders
    if remove_placeholders:
        for e in entries:
            if e.is_placeholder:
                remove_lines.add(e.line_number)
        placeholder_count = sum(1 for e in entries if e.is_placeholder)
        if placeholder_count:
            changes.append(f"Removed {placeholder_count} placeholder entries")
    else:
        placeholder_count = 0

    # 2. Remove duplicates (keep the longer entry, or the later one if same length)
    dup_count = 0
    if remove_duplicates:
        dup_groups = _find_duplicate_groups(entries)
        for keeper_idx, dup_idxs in dup_groups:
            # Among keeper + duplicates, keep the longest entry
            all_idxs = [keeper_idx] + dup_idxs
            best_idx = max(all_idxs, key=lambda i: len(entries[i].text))
            for idx in all_idxs:
                if idx != best_idx:
                    remove_lines.add(entries[idx].line_number)
                    dup_count += 1
        if dup_count:
            changes.append(f"Removed {dup_count} duplicate entries")

    # 3. Archive stale entries
    stale_count = 0
    if archive_stale:
        for e in entries:
            if e.line_number in remove_lines:
                continue  # already being removed
            if e.date_tag and not e.is_placeholder:
                try:
                    entry_date = datetime.strptime(e.date_tag, "%Y-%m-%d")
                    if entry_date < stale_cutoff:
                        remove_lines.add(e.line_number)
                        archive_entries.append(e)
                        stale_count += 1
                except ValueError:
                    pass
        if stale_count:
            changes.append(f"Archived {stale_count} stale entries (>{stale_days} days old)")

    # 4. Rebuild the file
    lines = content.split("\n")
    new_lines = []
    for i, line in enumerate(lines, 1):
        if i not in remove_lines:
            new_lines.append(line)

    # Remove consecutive blank lines (cleanup after removals)
    cleaned = []
    prev_blank = False
    for line in new_lines:
        is_blank = line.strip() == ""
        if is_blank and prev_blank:
            continue
        cleaned.append(line)
        prev_blank = is_blank

    new_content = "\n".join(cleaned)

    # Add archived section if we moved entries there
    if archive_entries:
        archive_section = "\n\n---\n\n## Archived (compacted)\n\n"
        archive_section += f"_Moved by `shelfai compact` on {today.strftime('%Y-%m-%d')}. "
        archive_section += f"Entries older than {stale_days} days._\n\n"
        for e in archive_entries:
            archive_section += f"- {e.text}\n"
        new_content = new_content.rstrip() + archive_section

    compacted_tokens = _estimate_tokens(new_content)
    compacted_lines = new_content.count("\n") + 1

    real_entries_original = sum(1 for e in entries if not e.is_placeholder)
    real_entries_compacted = real_entries_original - dup_count - stale_count

    if not changes:
        changes.append("No changes needed")

    return CompactResult(
        path=filepath,
        original_entries=real_entries_original,
        compacted_entries=real_entries_compacted,
        removed_duplicates=dup_count,
        removed_placeholders=placeholder_count,
        archived_stale=stale_count,
        original_lines=original_lines,
        compacted_lines=compacted_lines,
        original_tokens=original_tokens,
        compacted_tokens=compacted_tokens,
        content=new_content,
        changes=changes,
    )
