"""
Obsidian MCP Server — Expose Obsidian vaults as MCP tools.

Transport: StreamableHTTP (reachable cross-OS from WSL2 to Windows host).
Run: python server.py --vaults-root "D:\\Obsidian" --port 8765
"""

from __future__ import annotations

import argparse
import json
import os
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from fastmcp import FastMCP

# ── Config ────────────────────────────────────────────────────────────────

def resolve_vaults_root(cli_arg: str | None = None) -> Path:
    """Resolve vaults root: CLI arg > env var > config.json > cwd."""
    if cli_arg:
        return Path(cli_arg).expanduser().resolve()

    env_val = os.environ.get("OBSIDIAN_VAULTS_ROOT")
    if env_val:
        return Path(env_val).expanduser().resolve()

    config_path = Path(__file__).parent / "config.json"
    if config_path.exists():
        data = json.loads(config_path.read_text(encoding="utf-8"))
        if root := data.get("vaults_root"):
            return Path(root).expanduser().resolve()

    # Fallback: assume vaults/ is next to this script
    return (Path(__file__).parent / "vaults").resolve()


VAULTS_ROOT: Path = Path(".")  # set on startup
OBSIDIAN_HIDDEN = {".obsidian", ".trash", ".git", ".DS_Store", "node_modules"}
MARKDOWN_EXT = {".md", ".markdown"}
DEFAULT_EXCLUDE_EXT = {".webp", ".png", ".jpg", ".jpeg", ".gif", ".svg",
                       ".pdf", ".mp4", ".mp3", ".mov", ".excalidraw", ".bmp",
                       ".ico", ".tiff", ".heic", ".avif"}


# ── Helpers ───────────────────────────────────────────────────────────────

def _compact_date(timestamp: float) -> str:
    """Return YYYY-MM-DD (no time/microseconds/tz — 1/10 the token cost)."""
    return datetime.fromtimestamp(timestamp, tz=timezone.utc).strftime("%Y-%m-%d")


def _is_hidden_path(rel: Path) -> bool:
    """True if any path component is a dotfile/dotdir (.obsidian, .git, .trash, etc.)."""
    return any(part.startswith(".") for part in rel.parts)


# ── Helpers ───────────────────────────────────────────────────────────────

def _vault_path(vault_name: str) -> Path:
    """Get the absolute path to a vault directory, with safety check."""
    vault = (VAULTS_ROOT / vault_name).resolve()
    # Ensure it's actually inside VAULTS_ROOT (prevent traversal)
    if not str(vault).startswith(str(VAULTS_ROOT.resolve())):
        raise ValueError(f"Vault name escapes root: {vault_name}")
    if not vault.exists():
        raise FileNotFoundError(f"Vault not found: {vault_name}")
    if not vault.is_dir():
        raise NotADirectoryError(f"Not a directory: {vault_name}")
    return vault


def _resolve_doc(vault_name: str, doc_path: str) -> Path:
    """Resolve a document path within a vault, preventing traversal."""
    vault = _vault_path(vault_name)
    # Normalize: strip leading slash, handle backslashes
    clean = doc_path.lstrip("/\\").replace("\\", "/")
    target = (vault / clean).resolve()
    if not str(target).startswith(str(vault)):
        raise ValueError(f"Document path escapes vault: {doc_path}")
    return target


def _file_info_lean(p: Path, relative_to: Path) -> dict[str, Any]:
    """Bare index entry — name, path, type only. ~55 chars vs ~120."""
    return {
        "name": p.name,
        "path": str(p.relative_to(relative_to)).replace("\\", "/"),
        "type": "directory" if p.is_dir() else "file",
    }


def _file_info_full(p: Path, relative_to: Path) -> dict[str, Any]:
    """Full metadata for a single file — size, date, counts."""
    stat = p.stat()
    info: dict[str, Any] = {
        "name": p.name,
        "path": str(p.relative_to(relative_to)).replace("\\", "/"),
        "type": "directory" if p.is_dir() else "file",
        "size_bytes": stat.st_size if p.is_file() else None,
        "modified": _compact_date(stat.st_mtime),
    }
    if p.is_file() and _is_markdown(p):
        try:
            content = p.read_text(encoding="utf-8", errors="replace")
        except Exception:
            content = ""
        info["line_count"] = content.count("\n") + (1 if content else 0)
        info["char_count"] = len(content)
    return info


def _is_markdown(p: Path) -> bool:
    return p.suffix.lower() in MARKDOWN_EXT


def _search_file(filepath: Path, query: str) -> list[dict[str, Any]]:
    """Search a single markdown file for query, return match contexts."""
    try:
        text = filepath.read_text(encoding="utf-8", errors="replace")
    except Exception:
        return []
    results = []
    query_lower = query.lower()
    for i, line in enumerate(text.splitlines(), start=1):
        if query_lower in line.lower():
            # Trim context snippet
            idx = line.lower().index(query_lower)
            start = max(0, idx - 40)
            end = min(len(line), idx + len(query) + 40)
            snippet = line[start:end]
            if start > 0:
                snippet = "…" + snippet
            if end < len(line):
                snippet += "…"
            results.append({"line": i, "snippet": snippet.strip()})
    return results


def _parse_wikilinks(content: str) -> list[str]:
    """Extract [[wikilink]] targets from markdown content."""
    # Match [[Target]] or [[Target|Alias]]
    pattern = re.compile(r"\[\[([^\]|#]+?)(?:[|#][^\]]+)?\]\]")
    return [m.group(1).strip() for m in pattern.finditer(content)]


# ── Server ────────────────────────────────────────────────────────────────

mcp = FastMCP("Obsidian Vaults")


# ── Tools ─────────────────────────────────────────────────────────────────

@mcp.tool()
def get_vaults() -> dict[str, Any]:
    """FILESYSTEM TOOL (not memory) — List Obsidian vaults on disk. Use to discover available vaults before browsing one.

    Returns a list of vaults with name, path, and note counts.
    """
    if not VAULTS_ROOT.exists():
        return {"vaults": [], "vaults_root": str(VAULTS_ROOT), "error": "Vaults root does not exist."}

    vaults = []
    for entry in sorted(VAULTS_ROOT.iterdir()):
        if not entry.is_dir():
            continue
        if entry.name.startswith("."):
            continue
        # Check if it looks like a vault (has .obsidian or contains .md files)
        has_obsidian = (entry / ".obsidian").is_dir()
        md_count = sum(1 for f in entry.rglob("*.md") if ".obsidian" not in str(f) and ".trash" not in str(f))
        if has_obsidian or md_count > 0:
            vaults.append({
                "name": entry.name,
                "path": str(entry),
                "note_count": md_count,
                "has_obsidian_config": has_obsidian,
            })

    return {"vaults": vaults, "vaults_root": str(VAULTS_ROOT), "count": len(vaults)}


@mcp.tool()
def list_vault(
    vault_name: str,
    path: str = "",
    recursive: bool = False,
    include_files: bool = True,
    include_directories: bool = True,
    file_extensions: list[str] | None = None,  # None = all; pass ['.md'] for notes only
    exclude_extensions: list[str] | None = None,  # None = default image blacklist (see below)
    include_hidden: bool = False,
) -> dict[str, Any]:
    """FILESYSTEM TOOL (not memory) — Index files and folders in vault on disk. Use to see structure before reading docs.

    Returns a lean listing (name + path + type only). Image/binary files (.webp,
    .png, .jpg, .gif, .svg, .pdf, .mp4, .mp3, .mov, .excalidraw) are excluded by
    default — the model can't read them anyway. Pass exclude_extensions=[] to see
    everything. Use get_file_info for size/date/counts on a specific file.
    Hidden dirs are excluded by default.

    Args:
        vault_name: Name of the vault (subdirectory name under vaults root)
        path: Subfolder inside the vault. Omit or pass "" for vault root.
            NOT the vault name — that's vault_name. Examples:
            "" or omit = vault root, "NPCs" = vault/NPCs/, "NPCs/Hostile" = deeper.
        recursive: If True, list all nested contents recursively.
        include_files: Include files in results (default True).
        include_directories: Include directories in results (default True).
        file_extensions: Filter files by extension. Pass ['.md'] for notes only,
            ['.md', '.canvas'] for notes + canvases, or omit/None for all files.
        exclude_extensions: Extensions to EXCLUDE. Defaults to binary/image types
            (webp, png, jpg, gif, svg, pdf, mp4, mp3, mov, excalidraw).
            Pass [] to disable the blacklist.
        include_hidden: If True, include .obsidian config directory (default False).

    Usage:
        list_vault("obsidian-skt")                      → all files (images excluded) at root
        list_vault("obsidian-skt", file_extensions=['.md'])  → notes only
        list_vault("obsidian-skt", "", True, file_extensions=['.md'])  → all notes recursively
        list_vault("obsidian-skt", exclude_extensions=[])  → everything including images

    Returns:
        directories: [{name, path, type}]
        files: [{name, path, type}]
        total: combined count
    """
    vault = _vault_path(vault_name)
    # Normalize the inner path
    clean_path = path.lstrip("/\\").replace("\\", "/")
    target = (vault / clean_path).resolve() if clean_path else vault

    if not str(target).startswith(str(vault)):
        return {"error": "Path escapes vault boundary.", "vault": vault_name}

    if not target.exists():
        return {"error": f"Path not found: {path or '/'}", "vault": vault_name}

    # Resolve exclude list: explicit arg > default image blacklist
    if exclude_extensions is None:
        exclude_set = DEFAULT_EXCLUDE_EXT
    else:
        exclude_set = {e.lower() for e in exclude_extensions}

    directories = []
    files = []

    if recursive:
        iterator = target.rglob("*")
    else:
        iterator = target.iterdir()

    for entry in sorted(iterator, key=lambda p: (not p.is_dir(), p.name.lower())):
        rel = entry.relative_to(vault)

        # Filter hidden directories/files unless include_hidden is set
        if _is_hidden_path(rel) and not include_hidden:
            continue

        info = _file_info_lean(entry, vault)
        if entry.is_dir():
            if include_directories:
                directories.append(info)
        else:
            if not include_files:
                continue
            if file_extensions:
                # Case-insensitive extension check
                ext_lower = entry.suffix.lower()
                if not any(ext_lower == fe.lower() for fe in file_extensions):
                    continue
            if entry.suffix.lower() in exclude_set:
                continue
            files.append(info)

    return {
        "vault": vault_name,
        "path": str(target.relative_to(vault)).replace("\\", "/") or "/",
        "directories": directories,
        "files": files,
        "total": len(directories) + len(files),
    }


@mcp.tool()
def get_vault_document(vault_name: str, doc_path: str) -> dict[str, Any]:
    """FILESYSTEM TOOL (not memory) — Read a markdown note from vault on disk. Use after list_vault to get content.

    Args:
        vault_name: Name of the vault.
        doc_path: Relative path to the document within the vault (e.g. 'folder/note.md').

    Returns the document content with metadata.
    """
    target = _resolve_doc(vault_name, doc_path)
    if not target.exists():
        return {"error": f"Document not found: {doc_path}", "vault": vault_name}
    if not _is_markdown(target):
        return {"error": f"Not a markdown file: {doc_path}", "vault": vault_name}

    content = target.read_text(encoding="utf-8", errors="replace")
    stat = target.stat()
    lines = content.splitlines()

    return {
        "vault": vault_name,
        "path": doc_path.lstrip("/\\").replace("\\", "/"),
        "content": content,
        "line_count": len(lines),
        "word_count": len(content.split()),
        "char_count": len(content),
        "size_bytes": stat.st_size,
        "modified": _compact_date(stat.st_mtime),
    }


@mcp.tool()
def get_file_info(vault_name: str, doc_path: str) -> dict[str, Any]:
    """FILESYSTEM TOOL (not memory) — Get size, date, and line count for a single file. Use after list_vault to inspect a specific note.

    Args:
        vault_name: Name of the vault.
        doc_path: Relative path to the file within the vault (e.g. 'folder/note.md').

    Returns metadata dict with: name, path, type, size_bytes, modified (YYYY-MM-DD),
    and for markdown files: line_count, char_count.
    """
    target = _resolve_doc(vault_name, doc_path)
    if not target.exists():
        return {"error": f"File not found: {doc_path}", "vault": vault_name}
    return _file_info_full(target, _vault_path(vault_name))


@mcp.tool()
def search_vault(vault_name: str, query: str, max_results: int = 50) -> dict[str, Any]:
    """FILESYSTEM TOOL (not memory) — Search markdown files on disk. Find notes about characters, locations, concepts.

    Args:
        vault_name: Name of the vault to search.
        query: Search query (case-insensitive substring match).
        max_results: Maximum number of match results to return (default 50).

    Returns matches grouped by file with line numbers and context snippets.
    """
    vault = _vault_path(vault_name)
    results = []
    files_searched = 0

    for md_file in vault.rglob("*.md"):
        rel_path = md_file.relative_to(vault)
        if _is_hidden_path(rel_path):
            continue

        files_searched += 1
        matches = _search_file(md_file, query)
        if matches:
            results.append({
                "file": str(rel_path).replace("\\", "/"),
                "matches": matches[:max_results],
                "match_count": len(matches),
            })

        if len(results) >= max_results:
            break

    # Sort by match count descending
    results.sort(key=lambda r: r["match_count"], reverse=True)

    return {
        "vault": vault_name,
        "query": query,
        "results": results[:max_results],
        "total_matches": sum(r["match_count"] for r in results),
        "files_searched": files_searched,
        "files_with_matches": len(results),
    }


@mcp.tool()
def get_recent_changes(vault_name: str, limit: int = 20) -> dict[str, Any]:
    """Get recently modified documents in a vault.

    Args:
        vault_name: Name of the vault.
        limit: Maximum number of results (default 20).

    Returns documents sorted by modification time, most recent first.
    """
    vault = _vault_path(vault_name)
    md_files = []

    for md_file in vault.rglob("*.md"):
        rel_path = md_file.relative_to(vault)
        if _is_hidden_path(rel_path):
            continue
        rel = str(rel_path)
        stat = md_file.stat()
        md_files.append({
            "path": rel.replace("\\", "/"),
            "modified": _compact_date(stat.st_mtime),
        })

    md_files.sort(key=lambda f: f["modified"], reverse=True)

    return {
        "vault": vault_name,
        "recent": md_files[:limit],
        "total_notes": len(md_files),
    }


@mcp.tool()
def get_backlinks(vault_name: str, doc_path: str) -> dict[str, Any]:
    """Find all documents that link to the given document via [[wikilinks]].

    Args:
        vault_name: Name of the vault.
        doc_path: Relative path to the target document.

    Returns a list of files that link to it, with context snippets.
    """
    vault = _vault_path(vault_name)
    target = _resolve_doc(vault_name, doc_path)
    target_rel = str(target.relative_to(vault)).replace("\\", "/")
    target_stem = target.stem  # filename without extension

    backlinks = []

    for md_file in vault.rglob("*.md"):
        rel_path = md_file.relative_to(vault)
        if _is_hidden_path(rel_path):
            continue
        if md_file.resolve() == target.resolve():
            continue  # skip self

        try:
            content = md_file.read_text(encoding="utf-8", errors="replace")
        except Exception:
            continue

        # Check for wikilinks matching either the full path or stem
        pattern = re.compile(r"\[\[([^\]|#]+?)(?:[|#][^\]]+)?\]\]")
        matching_links = []
        for m in pattern.finditer(content):
            link_target = m.group(1).strip()
            # Match if link_target matches stem, filename, or full path
            if link_target == target_stem or link_target == target_rel or link_target == target.name:
                # Get line context
                line_no = content[:m.start()].count("\n") + 1
                lines = content.splitlines()
                line_text = lines[line_no - 1] if line_no <= len(lines) else ""
                matching_links.append({"line": line_no, "snippet": line_text.strip()[:120]})

        if matching_links:
            stat = md_file.stat()
            backlinks.append({
                "file": str(rel_path).replace("\\", "/"),
                "links": matching_links,
                "link_count": len(matching_links),
                "modified": _compact_date(stat.st_mtime),
            })

    backlinks.sort(key=lambda b: b["link_count"], reverse=True)

    return {
        "vault": vault_name,
        "target": doc_path.lstrip("/\\").replace("\\", "/"),
        "backlinks": backlinks,
        "total_backlinks": len(backlinks),
    }


@mcp.tool()
def create_document(vault_name: str, doc_path: str, content: str) -> dict[str, Any]:
    """Create a new markdown document in a vault.

    Args:
        vault_name: Name of the vault.
        doc_path: Relative path for the new document (must end in .md).
        content: Markdown content for the document.

    Returns the created document info.
    """
    if not doc_path.lower().endswith(".md"):
        doc_path = doc_path.rstrip("/\\") + ".md"

    target = _resolve_doc(vault_name, doc_path)
    if target.exists():
        return {"error": f"Document already exists: {doc_path}", "vault": vault_name}

    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(content, encoding="utf-8")

    return {
        "vault": vault_name,
        "path": doc_path.lstrip("/\\").replace("\\", "/"),
        "action": "created",
        "size_bytes": target.stat().st_size,
    }


@mcp.tool()
def edit_document(vault_name: str, doc_path: str, content: str) -> dict[str, Any]:
    """Edit (overwrite) an existing markdown document in a vault.

    Args:
        vault_name: Name of the vault.
        doc_path: Relative path to the document.
        content: New markdown content (replaces entire file).

    Returns the edited document info.
    """
    if not doc_path.lower().endswith(".md"):
        doc_path = doc_path.rstrip("/\\") + ".md"

    target = _resolve_doc(vault_name, doc_path)
    if not target.exists():
        return {"error": f"Document does not exist: {doc_path}", "vault": vault_name}

    target.write_text(content, encoding="utf-8")

    return {
        "vault": vault_name,
        "path": doc_path.lstrip("/\\").replace("\\", "/"),
        "action": "edited",
        "size_bytes": target.stat().st_size,
    }


@mcp.tool()
def create_vault_directory(vault_name: str, dir_path: str) -> dict[str, Any]:
    """FILESYSTEM TOOL (not memory) — Create a new directory in a vault on disk.

    Use before create_document when building out folder structures, or for
    organizing notes into subdirectories.

    Args:
        vault_name: Name of the vault.
        dir_path: Relative path for the new directory (e.g. 'Session-12' or 'NPCs/Dragon-Cult').

    Returns the created directory path.
    """
    # Normalize: strip leading slash, handle backslashes
    clean = dir_path.lstrip("/\\").replace("\\", "/")
    target = _resolve_doc(vault_name, clean)

    if target.exists():
        if not target.is_dir():
            return {"error": f"Path exists but is not a directory: {dir_path}", "vault": vault_name}
        return {"vault": vault_name, "path": str(target.relative_to(_vault_path(vault_name))).replace("\\", "/"),
                "action": "exists", "note": "Directory already exists"}

    target.mkdir(parents=True, exist_ok=True)

    return {
        "vault": vault_name,
        "path": str(target.relative_to(_vault_path(vault_name))).replace("\\", "/"),
        "action": "created",
    }


@mcp.tool()
def move_document(vault_name: str, src_path: str, dest_path: str) -> dict[str, Any]:
    """FILESYSTEM TOOL (not memory) — Move or rename a document within a vault.

    Moves a document from src_path to dest_path. Both paths are relative to the
    vault root. Automatically creates parent directories at the destination if
    they don't exist.

    Args:
        vault_name: Name of the vault.
        src_path: Current relative path of the document (e.g. 'Notes/draft.md').
        dest_path: Target relative path (e.g. 'Archive/draft.md' or 'Notes/final.md').

    Returns the move result with both old and new paths.
    """
    src = _resolve_doc(vault_name, src_path)
    if not src.exists():
        return {"error": f"Source document does not exist: {src_path}", "vault": vault_name}
    if not src.is_file():
        return {"error": f"Source is not a file: {src_path}", "vault": vault_name}

    dest = _resolve_doc(vault_name, dest_path)
    if dest.exists():
        return {"error": f"Destination already exists: {dest_path}", "vault": vault_name}

    dest.parent.mkdir(parents=True, exist_ok=True)
    src.rename(dest)

    vault = _vault_path(vault_name)
    return {
        "vault": vault_name,
        "old_path": src_path.lstrip("/\\").replace("\\", "/"),
        "new_path": str(dest.relative_to(vault)).replace("\\", "/"),
        "action": "moved",
    }


# ── Entrypoint ────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Obsidian MCP Server")
    parser.add_argument(
        "--vaults-root",
        default=None,
        help="Root directory containing Obsidian vaults (e.g. D:\\Obsidian)",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8765,
        help="Port for StreamableHTTP transport (default: 8765)",
    )
    parser.add_argument(
        "--host",
        default="0.0.0.0",
        help="Bind host (default: 0.0.0.0 for LAN access)",
    )
    args = parser.parse_args()

    global VAULTS_ROOT
    VAULTS_ROOT = resolve_vaults_root(args.vaults_root)

    print(f"Obsidian MCP Server")
    print(f"  Vaults root: {VAULTS_ROOT}")
    print(f"  Listening:   http://{args.host}:{args.port}/mcp")
    print(f"  Vaults found: {len(get_vaults()['vaults'])}")

    mcp.run(transport="streamable-http", host=args.host, port=args.port, path="/mcp")


if __name__ == "__main__":
    main()