"""MemGPT-style self-editable working context blocks.

Each block is a named, size-capped string the agent sees in its system
prompt on every turn and can edit via memory_insert / memory_replace /
memory_rethink tool calls.  Blocks are persisted across sessions.
"""

from __future__ import annotations

import json
import logging
import re
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, Optional

logger = logging.getLogger(__name__)

# Block name -> max character capacity
BLOCK_LIMITS: Dict[str, int] = {
    "project_context": 2000,
    "session_state": 1000,
    "user_preferences": 500,
    "active_findings": 500,
}

VALID_BLOCK_NAMES = frozenset(BLOCK_LIMITS.keys())

# Blocks that persist to memory.md across sessions
PERSISTENT_BLOCKS = frozenset({"project_context", "user_preferences"})

# Blocks that are archived to session JSONL at session end
SESSION_BLOCKS = frozenset({"session_state"})

# Blocks that are auto-populated (not user-edited)
AUTO_BLOCKS = frozenset({"active_findings"})

_BLOCK_MARKER_RE = re.compile(
    r"<!-- block:(\w+) -->\n(.*?)<!-- /block:\1 -->",
    re.DOTALL,
)


@dataclass
class MemoryBlocks:
    """Container for all working context blocks.

    Thread safety: pydantic-ai awaits tool calls sequentially within a
    single agent run, so no lock is needed.  The blocks dict is only
    mutated by tool calls and session lifecycle methods.
    """

    _blocks: Dict[str, str] = field(default_factory=dict)
    _session_id: str = field(default_factory=lambda: str(uuid.uuid4()))

    def __post_init__(self) -> None:
        for name in VALID_BLOCK_NAMES:
            if name not in self._blocks:
                self._blocks[name] = ""

    @property
    def session_id(self) -> str:
        return self._session_id

    def get(self, block: str) -> str:
        """Read a block's content."""
        if block not in VALID_BLOCK_NAMES:
            raise ValueError(
                f"Unknown block '{block}'. "
                f"Valid: {', '.join(sorted(VALID_BLOCK_NAMES))}"
            )
        return self._blocks.get(block, "")

    def insert(self, block: str, content: str) -> str:
        """Append content to a block.  Returns status message."""
        if block not in VALID_BLOCK_NAMES:
            return (
                f"Error: unknown block '{block}'. "
                f"Valid: {', '.join(sorted(VALID_BLOCK_NAMES))}"
            )
        if block in AUTO_BLOCKS:
            return f"Error: '{block}' is auto-populated and cannot be edited."

        limit = BLOCK_LIMITS[block]
        current = self._blocks.get(block, "")
        new = (current.rstrip() + "\n" + content).strip() if current.strip() else content.strip()

        if len(new) > limit:
            return (
                f"Error: insert would exceed {block} limit "
                f"({len(new)}/{limit} chars). "
                f"Use memory_replace to edit existing content, "
                f"or memory_rethink to rewrite."
            )

        self._blocks[block] = new
        return f"OK. {block} updated ({len(new)}/{limit} chars)."

    def replace(self, block: str, old_content: str, new_content: str) -> str:
        """Replace specific text within a block.  Returns status message."""
        if block not in VALID_BLOCK_NAMES:
            return (
                f"Error: unknown block '{block}'. "
                f"Valid: {', '.join(sorted(VALID_BLOCK_NAMES))}"
            )
        if block in AUTO_BLOCKS:
            return f"Error: '{block}' is auto-populated and cannot be edited."

        current = self._blocks.get(block, "")
        if old_content not in current:
            preview = current[:200] if current else "[empty]"
            return (
                f"Error: old_content not found in '{block}'. "
                f"Current content:\n{preview}"
            )

        updated = current.replace(old_content, new_content, 1)
        limit = BLOCK_LIMITS[block]
        if len(updated) > limit:
            return (
                f"Error: replace would exceed {block} limit "
                f"({len(updated)}/{limit} chars)."
            )

        self._blocks[block] = updated
        return f"OK. {block} updated ({len(updated)}/{limit} chars)."

    def rethink(self, block: str, new_content: str) -> str:
        """Complete rewrite of a block.  Returns status message."""
        if block not in VALID_BLOCK_NAMES:
            return (
                f"Error: unknown block '{block}'. "
                f"Valid: {', '.join(sorted(VALID_BLOCK_NAMES))}"
            )
        if block in AUTO_BLOCKS:
            return f"Error: '{block}' is auto-populated and cannot be edited."

        limit = BLOCK_LIMITS[block]
        new_content = new_content.strip()
        if len(new_content) > limit:
            return (
                f"Error: content exceeds {block} limit "
                f"({len(new_content)}/{limit} chars)."
            )

        self._blocks[block] = new_content
        return f"OK. {block} rewritten ({len(new_content)}/{limit} chars)."

    def render_for_prompt(self) -> str:
        """Render all blocks as a system prompt section."""
        lines = [
            "## Working Context (self-editable)\n",
            "You can read and edit these blocks with "
            "memory_insert, memory_replace, and memory_rethink.\n",
        ]
        for name in ["project_context", "session_state",
                      "user_preferences", "active_findings"]:
            content = self._blocks.get(name, "")
            limit = BLOCK_LIMITS[name]
            display = content if content else "[empty]"
            lines.append(f"### {name} ({len(content)}/{limit} chars)")
            lines.append(display)
            lines.append("")
        return "\n".join(lines)

    def total_chars(self) -> int:
        """Total characters across all blocks."""
        return sum(len(v) for v in self._blocks.values())


# ---------------------------------------------------------------------------
# Session lifecycle
# ---------------------------------------------------------------------------

def load_blocks_from_disk(project_path: str) -> MemoryBlocks:
    """Load memory blocks at session start.

    1. Load project_context + user_preferences from .hikaflow/memory.md
       (structured markers, or fallback: entire content -> project_context)
    2. session_state starts empty
    3. Auto-populate active_findings from .hikaflow/findings.json
    """
    blocks = MemoryBlocks()
    hikaflow_dir = Path(project_path) / ".hikaflow"

    # 1. Load persistent blocks from memory.md
    memory_path = hikaflow_dir / "memory.md"
    if memory_path.is_file():
        try:
            raw = memory_path.read_text(encoding="utf-8").strip()
            parsed = _parse_structured_memory(raw)
            if parsed:
                for bname in PERSISTENT_BLOCKS:
                    if bname in parsed:
                        limit = BLOCK_LIMITS[bname]
                        blocks._blocks[bname] = parsed[bname][:limit]
            elif raw:
                # Legacy freeform: treat entire content as project_context
                limit = BLOCK_LIMITS["project_context"]
                blocks._blocks["project_context"] = raw[:limit]
        except Exception:
            logger.warning("Failed to read memory.md", exc_info=True)

    # 2. session_state starts empty (already initialized)

    # 3. Auto-populate active_findings
    findings_path = hikaflow_dir / "findings.json"
    if findings_path.is_file():
        try:
            findings_data = json.loads(
                findings_path.read_text(encoding="utf-8")
            )
            if isinstance(findings_data, list):
                blocks._blocks["active_findings"] = _summarize_findings(
                    findings_data, BLOCK_LIMITS["active_findings"]
                )
        except Exception:
            logger.warning("Failed to read findings.json", exc_info=True)

    return blocks


def persist_blocks_to_disk(blocks: MemoryBlocks, project_path: str) -> None:
    """Persist memory blocks at session end.

    1. Write project_context + user_preferences to .hikaflow/memory.md
    2. Archive session_state to .hikaflow/sessions/{session_id}.jsonl
    """
    hikaflow_dir = Path(project_path) / ".hikaflow"
    hikaflow_dir.mkdir(parents=True, exist_ok=True)

    # 1. Persist project_context + user_preferences (V-08: with file locking)
    memory_path = hikaflow_dir / "memory.md"
    sections = []
    for bname in ["project_context", "user_preferences"]:
        content = blocks.get(bname)
        if content.strip():
            sections.append(
                f"<!-- block:{bname} -->\n{content}\n<!-- /block:{bname} -->"
            )
    try:
        import fcntl
        with open(memory_path, "w", encoding="utf-8") as f:
            fcntl.flock(f.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
            try:
                f.write("\n\n".join(sections))
            finally:
                fcntl.flock(f.fileno(), fcntl.LOCK_UN)
    except (OSError, BlockingIOError):
        # Another session holds the lock — skip rather than corrupt
        logger.warning("Skipped writing memory.md — file locked by another session")
    except Exception:
        logger.warning("Failed to write memory.md", exc_info=True)

    # 2. Archive session_state
    session_state = blocks.get("session_state")
    if session_state.strip():
        sessions_dir = hikaflow_dir / "sessions"
        sessions_dir.mkdir(parents=True, exist_ok=True)
        session_file = sessions_dir / f"{blocks.session_id}.jsonl"
        try:
            entry = {
                "session_id": blocks.session_id,
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "session_state": session_state,
            }
            with open(session_file, "a", encoding="utf-8") as f:
                f.write(json.dumps(entry) + "\n")
        except Exception:
            logger.warning("Failed to archive session state", exc_info=True)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _parse_structured_memory(raw: str) -> Optional[Dict[str, str]]:
    """Parse structured memory.md with ``<!-- block:name -->`` markers.

    Returns dict of block_name -> content, or None if no markers found.
    """
    result: Dict[str, str] = {}
    for match in _BLOCK_MARKER_RE.finditer(raw):
        block_name = match.group(1)
        content = match.group(2).strip()
        if block_name in VALID_BLOCK_NAMES:
            result[block_name] = content
    return result if result else None


def _summarize_findings(findings: list, max_chars: int) -> str:
    """Summarize open findings into a compact string."""
    open_findings = [
        f for f in findings
        if isinstance(f, dict) and f.get("status") != "fixed"
    ]
    if not open_findings:
        return ""

    lines = []
    for f in open_findings[:10]:
        sev = f.get("severity", "?")
        issue_id = f.get("issue_id", "?")
        file_ = f.get("file", "?")
        line_ = f.get("line", "?")
        desc = f.get("description", "")[:60]
        lines.append(f"[{sev}] {issue_id}: {file_}:{line_} - {desc}")

    summary = f"{len(open_findings)} open issues:\n" + "\n".join(lines)
    if len(summary) > max_chars:
        truncation_suffix = "\n... [truncated]"
        summary = summary[: max_chars - len(truncation_suffix)] + truncation_suffix
    return summary
