"""
MCP server — thin wrapper around the same logic the CLI uses.
Start with: ctxbaton serve
"""

from __future__ import annotations

import json
import os
from typing import Optional

from mcp.server.fastmcp import FastMCP

from . import db
from .snapshots import save_snapshot, list_snapshots, get_snapshot

_project_path = os.environ.get("MEMORY_PROJECT", os.getcwd())
db.init(_project_path)

mcp = FastMCP("ctxbaton")


@mcp.tool()
def get_context() -> dict:
    """Full project context snapshot for resuming work cold."""
    from .snapshots import last_snapshot
    spec      = db.list_memories(type="spec",    limit=1)
    todos     = db.list_memories(type="todo",    status="pending",     limit=20)
    todos    += db.list_memories(type="todo",    status="in_progress", limit=20)
    handoff   = db.list_memories(type="handoff", limit=1)
    decisions = db.list_memories(type="decision", limit=5)
    snap      = last_snapshot()
    return {
        "spec":      spec[0]    if spec    else None,
        "todos":     todos,
        "handoff":   handoff[0] if handoff else None,
        "decisions": decisions,
        "snapshot":  {k: snap[k] for k in ("id","timestamp","git_hash","git_branch") if snap} if snap else None,
    }


@mcp.tool()
def replay(snapshot_id: Optional[str] = None) -> dict:
    """
    Reconstruct what happened since a snapshot — git diff, todo changes,
    memories added. Call with no args to replay from the last snapshot.
    """
    import subprocess
    from .snapshots import _read_all, _current_todos, _compute_diff

    snaps = _read_all()
    if not snaps:
        return {"error": "No snapshots found"}

    anchor = get_snapshot(snapshot_id) if snapshot_id else snaps[-1]
    if not anchor:
        return {"error": f"Snapshot {snapshot_id} not found"}

    cwd        = anchor.get("cwd", os.getcwd())
    since_hash = anchor.get("git_hash")

    git_diff = None
    if since_hash:
        try:
            r = subprocess.run(
                ["git", "diff", "--stat", f"{since_hash}..HEAD"],
                cwd=cwd, capture_output=True, text=True, timeout=10
            )
            git_diff = r.stdout.strip() if r.returncode == 0 else None
        except Exception:
            pass

    todo_diff = _compute_diff(
        anchor.get("todos_state", []),
        _current_todos(),
    )

    memories_since = [
        m for m in db.list_memories(limit=50)
        if m["created_at"] > anchor["timestamp"]
    ]

    return {
        "anchor":          {"id": anchor["id"], "timestamp": anchor["timestamp"], "git_hash": since_hash},
        "git_diff_stat":   git_diff,
        "todo_changes":    todo_diff,
        "memories_added":  memories_since,
        "snapshots_since": len([s for s in snaps if s["timestamp"] > anchor["timestamp"]]),
    }


@mcp.tool()
def search_memories(
    query: str,
    type: Optional[str] = None,
    limit: int = 20,
) -> list[dict]:
    """Full-text search across memories."""
    return db.search_memories_fts(query, type=type, limit=limit)


@mcp.tool()
def save_memory(
    content: str,
    type: str = "general",
    summary: Optional[str] = None,
    status: Optional[str] = None,
    tags: Optional[list[str]] = None,
    agent: Optional[str] = None,
    session_id: Optional[str] = None,
) -> dict:
    """Save a structured memory to the current project."""
    import uuid
    if type not in db.MEMORY_TYPES:
        return {"error": f"Invalid type. Must be one of: {', '.join(sorted(db.MEMORY_TYPES))}"}
    if type == "todo" and not status:
        status = "pending"
    return db.insert_memory(
        id=str(uuid.uuid4()), content=content, summary=summary, type=type,
        status=status, tags=tags or [], agent=agent, session_id=session_id,
    )


@mcp.tool()
def list_memories(
    type: Optional[str] = None,
    status: Optional[str] = None,
    limit: int = 30,
) -> list[dict]:
    """List memories with optional filters."""
    return db.list_memories(type=type, status=status, limit=limit)


@mcp.tool()
def update_todo_status(id: str, status: str) -> dict:
    """Update todo status: pending | in_progress | done"""
    if status not in db.TODO_STATUSES:
        return {"error": f"Invalid status. Must be one of: {', '.join(sorted(db.TODO_STATUSES))}"}
    return db.update_memory(id=id, status=status) or {"error": f"Memory {id} not found"}


@mcp.tool()
def delete_memory(id: str) -> dict:
    """Delete a memory by ID."""
    return {"deleted": db.delete_memory(id), "id": id}


@mcp.tool()
def get_memory(id: str) -> dict:
    """Get a single memory by ID (or unique ID prefix)."""
    m = db.get_memory(id)
    if not m:
        all_m = db.list_memories(limit=1000)
        m = next((x for x in all_m if x["id"].startswith(id)), None)
    return m or {"error": f"Memory {id} not found"}


@mcp.tool()
def update_memory(
    id:      str,
    content: Optional[str]       = None,
    summary: Optional[str]       = None,
    type:    Optional[str]       = None,
    status:  Optional[str]       = None,
    tags:    Optional[list[str]] = None,
) -> dict:
    """Update a memory's content, type, status, or tags."""
    result = db.update_memory(id=id, content=content, summary=summary, type=type, status=status, tags=tags)
    return result or {"error": f"Memory {id} not found"}


@mcp.tool()
def resume() -> dict:
    """
    Full resumption context for starting a new agent session.
    Returns spec, todos, handoff, decisions, last snapshot, replay delta, and message count.
    """
    import subprocess
    from .snapshots import last_snapshot, _current_todos, _compute_diff

    spec      = db.list_memories(type="spec",     limit=1)
    todos     = db.list_memories(type="todo",     status="pending",     limit=20)
    todos    += db.list_memories(type="todo",     status="in_progress", limit=20)
    handoff   = db.list_memories(type="handoff",  limit=1)
    decisions = db.list_memories(type="decision", limit=5)
    snap      = last_snapshot()

    git_diff  = None
    todo_diff: list = []
    mem_since: list = []
    if snap:
        cwd        = snap.get("cwd", os.getcwd())
        since_hash = snap.get("git_hash")
        if since_hash:
            try:
                r = subprocess.run(
                    ["git", "diff", "--stat", f"{since_hash}..HEAD"],
                    cwd=cwd, capture_output=True, text=True, timeout=10,
                )
                git_diff = r.stdout.strip() if r.returncode == 0 else None
            except Exception:
                pass
        todo_diff = _compute_diff(snap.get("todos_state", []), _current_todos())
        mem_since = [m for m in db.list_memories(limit=20) if m["created_at"] > snap["timestamp"]]

    return {
        "spec":          spec[0]    if spec    else None,
        "todos":         todos,
        "handoff":       handoff[0] if handoff else None,
        "decisions":     decisions,
        "snapshot":      {k: snap[k] for k in ("id","timestamp","git_hash","git_branch","trigger") if snap and k in snap} if snap else None,
        "replay":        {"git_diff_stat": git_diff, "todo_changes": todo_diff, "memories_added": mem_since},
        "message_count": db.count_messages(),
    }


@mcp.tool()
def capture_session(
    agent:           str           = "unknown",
    session_id:      Optional[str] = None,
    cwd:             Optional[str] = None,
    transcript_path: Optional[str] = None,
    trigger:         str           = "manual",
) -> dict:
    """
    Capture current session state from any agent.
    Use trigger='handoff' when switching agents, 'stop' for normal end-of-turn.
    """
    from .snapshots import save_snapshot

    effective_cwd = cwd or _project_path
    new_messages  = 0
    if transcript_path:
        try:
            from .transcript import process_transcript
            offset               = db.get_transcript_cursor(transcript_path)
            inserted, new_offset = process_transcript(transcript_path, session_id=session_id, offset=offset)
            db.set_transcript_cursor(transcript_path, new_offset)
            new_messages = inserted
        except Exception:
            pass

    if session_id:
        db.upsert_session(
            id=session_id, agent=agent, project=effective_cwd,
            transcript_path=transcript_path,
        )

    snap = save_snapshot(
        session_id=session_id, agent=agent,
        transcript_path=transcript_path, cwd=effective_cwd,
        trigger=trigger,
    )
    return {
        "snapshot_id":      snap["id"],
        "timestamp":        snap["timestamp"],
        "messages_indexed": new_messages,
        "trigger":          trigger,
    }


@mcp.tool()
def save_handoff(
    message: str,
    todo_id: Optional[str] = None,
) -> dict:
    """
    Save a handoff note and capture session state.
    Call before switching agents or closing a session mid-task.
    """
    import uuid
    from .snapshots import save_snapshot

    m = db.insert_memory(
        id=str(uuid.uuid4()), content=message, type="handoff",
    )
    if todo_id:
        db.update_memory(id=todo_id, status="in_progress")

    snap = save_snapshot(trigger="handoff")
    return {"handoff_id": m["id"], "snapshot_id": snap["id"]}


@mcp.tool()
def search_messages(
    query:      str,
    session_id: Optional[str] = None,
    limit:      int = 20,
) -> list[dict]:
    """Full-text search over indexed transcript messages."""
    return db.search_messages_text(query, session_id=session_id, limit=limit)
