# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# Licensed under the MIT License

"""Filesystem structure graph — directory tree with size analysis.

Scans a directory and builds a hierarchical graph suitable for
force-directed or treemap visualization in the dashboard.
"""

from __future__ import annotations

import logging
import os
from datetime import datetime
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

# File type categories
EXT_CATEGORIES: dict[str, str] = {
    # Code
    ".py": "code", ".js": ".code", ".ts": "code", ".svelte": "code",
    ".html": "code", ".css": "code", ".json": "code", ".yaml": "code",
    ".yml": "code", ".toml": "code", ".sh": "code", ".bash": "code",
    ".rs": "code", ".go": "code", ".java": "code", ".c": "code",
    ".cpp": "code", ".h": "code", ".rb": "code", ".php": "code",
    ".sql": "code", ".xml": "code", ".md": "docs", ".txt": "docs",
    ".rst": "docs", ".csv": "data",
    # Media
    ".png": "image", ".jpg": "image", ".jpeg": "image", ".gif": "image",
    ".svg": "image", ".ico": "image", ".webp": "image", ".bmp": "image",
    ".mp3": "audio", ".wav": "audio", ".flac": "audio", ".ogg": "audio",
    ".mp4": "video", ".mkv": "video", ".avi": "video", ".mov": "video",
    # Data
    ".db": "database", ".sqlite": "database", ".sqlite3": "database",
    ".zip": "archive", ".tar": "archive", ".gz": "archive",
    ".bz2": "archive", ".xz": "archive", ".7z": "archive",
    ".pdf": "document", ".docx": "document", ".xlsx": "document",
    # Config
    ".env": "config", ".ini": "config", ".cfg": "config",
    ".conf": "config", ".lock": "config",
}

CATEGORY_COLORS: dict[str, str] = {
    "code": "#22c55e", "docs": "#0ea5e9", "data": "#f59e0b",
    "image": "#ec4899", "audio": "#a855f7", "video": "#8b5cf6",
    "database": "#ef4444", "archive": "#64748b", "document": "#06b6d4",
    "config": "#84cc16", "directory": "#475569", "other": "#334155",
}


def _get_category(path: Path) -> str:
    if path.is_dir():
        return "directory"
    return EXT_CATEGORIES.get(path.suffix.lower(), "other")


def _format_size(size_bytes: int) -> str:
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024 * 1024:
        return f"{size_bytes / 1024:.1f} KB"
    elif size_bytes < 1024 * 1024 * 1024:
        return f"{size_bytes / (1024 * 1024):.1f} MB"
    return f"{size_bytes / (1024 * 1024 * 1024):.2f} GB"


def scan_directory(
    root: str = "~",
    max_depth: int = 3,
    max_nodes: int = 1000,
    min_size: int = 0,
) -> dict[str, Any]:
    """Scan a directory tree and return a graph for visualization.

    Returns:
        {
            "nodes": [{"id": "/path", "name": "file.py", "size": 1234, "category": "code", ...}],
            "children": {"/parent": ["/child1", "/child2"]},
            "stats": {"total_files": N, "total_dirs": N, "total_size": N, ...}
        }
    """
    root_path = Path(root).expanduser().resolve()
    if not root_path.exists():
        return {"nodes": [], "children": {}, "stats": {"error": f"Path not found: {root}"}}

    nodes: list[dict[str, Any]] = []
    children: dict[str, list[str]] = {}
    total_files = 0
    total_dirs = 0
    total_size = 0
    category_sizes: dict[str, int] = {}
    category_counts: dict[str, int] = {}
    node_count = 0

    # Skip patterns
    skip_dirs = {
        "__pycache__", "node_modules", ".git", ".svelte-kit",
        ".cache", ".local", ".config", "snap", ".npm", ".nvm",
    }

    def _scan(path: Path, depth: int, parent_id: str | None) -> int:
        nonlocal total_files, total_dirs, total_size, node_count

        if node_count >= max_nodes:
            return 0
        if depth > max_depth:
            return 0

        node_id = str(path)
        is_dir = path.is_dir()

        if is_dir and path.name in skip_dirs:
            return 0

        try:
            if is_dir:
                total_dirs += 1
                dir_size = 0
                child_ids: list[str] = []

                try:
                    entries = sorted(path.iterdir(), key=lambda p: (not p.is_dir(), p.name.lower()))
                except PermissionError:
                    return 0

                for entry in entries:
                    child_size = _scan(entry, depth + 1, node_id)
                    if child_size > 0:
                        child_ids.append(str(entry))
                    dir_size += child_size

                if dir_size < min_size and depth > 0:
                    return dir_size  # skip small dirs but still count their size

                node_count += 1
                category = "directory"

                # Directory security: check if it contains keys or encrypted files
                dir_security = "none"
                dir_name = path.name.lower()
                if dir_name == "keys":
                    dir_security = "key"
                elif any(f.suffix == ".enc" for f in path.iterdir() if f.is_file()):
                    dir_security = "encrypted_rest"
                elif dir_name in (".ssh", ".gnupg", ".pki"):
                    dir_security = "secure"

                nodes.append({
                    "id": node_id,
                    "name": path.name or str(path),
                    "size": dir_size,
                    "size_fmt": _format_size(dir_size),
                    "category": category,
                    "is_dir": True,
                    "depth": depth,
                    "children_count": len(child_ids),
                    "modified": datetime.fromtimestamp(path.stat().st_mtime).isoformat() if path.exists() else None,
                    "security": dir_security,
                })
                if child_ids:
                    children[node_id] = child_ids
                if parent_id:
                    children.setdefault(parent_id, [])

                category_sizes["directory"] = category_sizes.get("directory", 0) + dir_size
                category_counts["directory"] = category_counts.get("directory", 0) + 1

                return dir_size
            else:
                try:
                    file_size = path.stat().st_size
                except (PermissionError, OSError):
                    return 0

                total_files += 1
                total_size += file_size

                if file_size < min_size:
                    return file_size  # count size but don't create node

                node_count += 1
                category = _get_category(path)

                # Security state detection
                security = "none"
                suffix = path.suffix.lower()
                if suffix == ".enc":
                    security = "encrypted_rest"  # encrypted at rest
                elif suffix == ".key":
                    security = "key"  # encryption key
                elif path.name.startswith(".") and suffix in (".env", ""):
                    security = "sensitive"  # dotfiles / env vars
                else:
                    # Check if an .enc companion exists
                    enc_companion = path.with_suffix(path.suffix + ".enc")
                    if enc_companion.exists():
                        security = "encrypted_rest"

                nodes.append({
                    "id": node_id,
                    "name": path.name,
                    "size": file_size,
                    "size_fmt": _format_size(file_size),
                    "category": category,
                    "is_dir": False,
                    "depth": depth,
                    "ext": suffix,
                    "modified": datetime.fromtimestamp(path.stat().st_mtime).isoformat(),
                    "security": security,
                })

                category_sizes[category] = category_sizes.get(category, 0) + file_size
                category_counts[category] = category_counts.get(category, 0) + 1

                return file_size
        except (PermissionError, OSError):
            return 0

    _scan(root_path, 0, None)

    return {
        "root": str(root_path),
        "nodes": nodes,
        "children": children,
        "stats": {
            "total_files": total_files,
            "total_dirs": total_dirs,
            "total_size": total_size,
            "total_size_fmt": _format_size(total_size),
            "total_nodes": len(nodes),
            "category_sizes": {k: {"size": v, "size_fmt": _format_size(v), "count": category_counts.get(k, 0)}
                               for k, v in sorted(category_sizes.items(), key=lambda x: -x[1])},
        },
        "colors": CATEGORY_COLORS,
    }
