"""Workspace context builder — scans directories for datasets, training outputs, GPU info.

Automatically detects what the user has in their current workspace so the
assistant can give contextually relevant advice without asking.
"""

from __future__ import annotations

import logging
import subprocess
import sys
from pathlib import Path

logger = logging.getLogger(__name__)


def build_workspace_context(path: str = ".") -> dict:
    """Scan workspace and build rich context for the system prompt.

    Returns:
        Dict with datasets, training_outputs, config_files, environment.
    """
    root = Path(path).resolve()
    context = {
        "datasets": [],
        "training_outputs": [],
        "config_files": [],
        "environment": _detect_environment(),
    }

    # Find datasets — cap file walking to avoid hanging on huge dirs
    _scan_datasets(root, context["datasets"])
    _scan_training_outputs(root, context["training_outputs"])
    _scan_configs(root, context["config_files"])

    return context


# ---------------------------------------------------------------------------
# Dataset detection
# ---------------------------------------------------------------------------


def _scan_datasets(root: Path, out: list[dict], max_depth: int = 4) -> None:
    """Detect datasets in the directory tree."""
    seen_paths: set[str] = set()

    # LeRobot datasets (v2/v3): have meta/info.json
    for info_json in _rglob_limited(root, "meta/info.json", max_depth):
        ds_path = str(info_json.parent.parent)
        if ds_path not in seen_paths:
            seen_paths.add(ds_path)
            out.append({
                "path": ds_path,
                "format": "lerobot",
                "version": _detect_lerobot_version(info_json),
            })

    # HDF5 files
    for hdf5 in _rglob_limited(root, "*.hdf5", max_depth):
        p = str(hdf5)
        if p not in seen_paths:
            seen_paths.add(p)
            out.append({"path": p, "format": "hdf5"})

    # MCAP files (ROS 2 bags)
    for mcap in _rglob_limited(root, "*.mcap", max_depth):
        p = str(mcap)
        if p not in seen_paths:
            seen_paths.add(p)
            out.append({"path": p, "format": "mcap"})

    # RLDS (TFRecord)
    for tfr in _rglob_limited(root, "*.tfrecord", max_depth):
        ds_dir = str(tfr.parent)
        if ds_dir not in seen_paths:
            seen_paths.add(ds_dir)
            out.append({"path": ds_dir, "format": "rlds"})

    # ROS 1 bags
    for bag in _rglob_limited(root, "*.bag", max_depth):
        p = str(bag)
        if p not in seen_paths:
            seen_paths.add(p)
            out.append({"path": p, "format": "rosbag"})

    # Parquet (possible standalone LeRobot data)
    for pq in _rglob_limited(root, "*.parquet", max_depth):
        ds_dir = str(pq.parent)
        if ds_dir not in seen_paths and not any(
            ds_dir.startswith(d["path"]) for d in out
        ):
            seen_paths.add(ds_dir)
            out.append({"path": ds_dir, "format": "parquet"})


def _detect_lerobot_version(info_json: Path) -> str:
    """Detect LeRobot dataset version from meta/info.json."""
    try:
        import json

        data = json.loads(info_json.read_text())
        codebase_version = data.get("codebase_version", "")
        if codebase_version.startswith("v2"):
            return "v2"
        if codebase_version.startswith("v3"):
            return "v3"
        # Heuristic: v2 uses "video_keys", v3 uses "features"
        if "features" in data:
            return "v3"
        if "video_keys" in data:
            return "v2"
        return "unknown"
    except Exception:
        return "unknown"


# ---------------------------------------------------------------------------
# Training output detection
# ---------------------------------------------------------------------------


def _scan_training_outputs(root: Path, out: list[dict], max_depth: int = 4) -> None:
    """Detect training outputs (TensorBoard events, logs, checkpoints)."""
    seen: set[str] = set()

    # TensorBoard events
    for ev in _rglob_limited(root, "events.out.tfevents.*", max_depth):
        p = str(ev.parent)
        if p not in seen:
            seen.add(p)
            out.append({"path": p, "type": "tensorboard"})

    # Text logs
    for log in _rglob_limited(root, "train.log", max_depth):
        p = str(log.parent)
        if p not in seen:
            seen.add(p)
            out.append({"path": p, "type": "text_log"})

    # Checkpoint directories
    for ckpt in _rglob_limited(root, "checkpoint-*", max_depth):
        if ckpt.is_dir():
            p = str(ckpt.parent)
            if p not in seen:
                seen.add(p)
                out.append({"path": p, "type": "checkpoints"})


# ---------------------------------------------------------------------------
# Config detection
# ---------------------------------------------------------------------------


def _scan_configs(root: Path, out: list[str]) -> None:
    """Find relevant config files (top level only to avoid noise)."""
    for pattern in ["*.yaml", "*.yml", "*.json", "*config*.py"]:
        for f in sorted(root.glob(pattern))[:10]:
            name_lower = f.name.lower()
            # Skip lockfiles and package manifests
            if name_lower in ("package.json", "package-lock.json", "tsconfig.json"):
                continue
            if any(kw in name_lower for kw in ("config", "modality", "train", "policy")):
                out.append(str(f))
            elif f.suffix in (".yaml", ".yml"):
                out.append(str(f))


# ---------------------------------------------------------------------------
# Environment detection
# ---------------------------------------------------------------------------


def _detect_environment() -> dict:
    """Detect Python version, GPU, installed packages."""
    return {
        "python_version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
        "gpu": _detect_gpu(),
        "installed_packages": _get_relevant_packages(),
        "ros_version": _detect_ros(),
    }


def _detect_gpu() -> dict:
    """Detect GPU type and VRAM via nvidia-smi."""
    try:
        result = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=name,memory.total,memory.free,driver_version",
                "--format=csv,noheader,nounits",
            ],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0 and result.stdout.strip():
            gpus = []
            for line in result.stdout.strip().split("\n"):
                parts = [p.strip() for p in line.split(",")]
                if len(parts) >= 4:
                    gpus.append({
                        "name": parts[0],
                        "vram_total_mb": int(float(parts[1])),
                        "vram_free_mb": int(float(parts[2])),
                        "driver": parts[3],
                    })
            if gpus:
                return {"available": True, "gpus": gpus}
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    # Check for Apple Silicon MPS
    try:
        result = subprocess.run(
            ["sysctl", "-n", "machdep.cpu.brand_string"],
            capture_output=True,
            text=True,
            timeout=3,
        )
        if result.returncode == 0 and "Apple" in result.stdout:
            return {"available": True, "gpus": [{"name": "Apple Silicon (MPS)", "vram_total_mb": 0, "vram_free_mb": 0, "driver": "MPS"}]}
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    return {"available": False, "gpus": []}


def _get_relevant_packages() -> dict[str, str]:
    """Check which robotics packages are installed."""
    packages: dict[str, str] = {}
    check_list = [
        "lerobot", "torch", "tensorflow", "jax", "mujoco",
        "h5py", "mcap", "openvla", "groot", "osmo",
        "orbit-robotics", "transformers", "opencv-python",
    ]
    for pkg in check_list:
        try:
            # Use importlib.metadata — faster than pip show
            from importlib.metadata import version as pkg_version

            packages[pkg] = pkg_version(pkg)
        except Exception:
            pass
    return packages


def _detect_ros() -> str | None:
    """Detect ROS 1 or 2 installation."""
    import os

    if os.environ.get("ROS_DISTRO"):
        ros_version = os.environ.get("ROS_VERSION", "?")
        distro = os.environ.get("ROS_DISTRO", "unknown")
        return f"ROS {ros_version} ({distro})"
    return None


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _rglob_limited(root: Path, pattern: str, max_depth: int) -> list[Path]:
    """Like Path.rglob but limited depth to avoid scanning huge trees."""
    results: list[Path] = []
    # Use iterative approach with depth tracking
    try:
        for match in root.rglob(pattern):
            # Check depth relative to root
            try:
                rel = match.relative_to(root)
                if len(rel.parts) <= max_depth + 1:
                    results.append(match)
                    if len(results) >= 50:  # safety cap
                        break
            except ValueError:
                continue
    except PermissionError:
        pass
    return results


def format_context_for_prompt(context: dict) -> str:
    """Format workspace context as readable text for the system prompt."""
    parts: list[str] = []

    if context["datasets"]:
        parts.append("Detected datasets:")
        for ds in context["datasets"]:
            ver = f" ({ds['version']})" if ds.get("version") else ""
            parts.append(f"  - {ds['path']} [{ds['format']}{ver}]")
    else:
        parts.append("No datasets detected in current directory.")

    if context["training_outputs"]:
        parts.append("\nTraining outputs found:")
        for to in context["training_outputs"]:
            parts.append(f"  - {to['path']} [{to['type']}]")

    if context["config_files"]:
        parts.append("\nConfig files:")
        for cf in context["config_files"]:
            parts.append(f"  - {cf}")

    env = context.get("environment", {})
    parts.append(f"\nPython: {env.get('python_version', 'unknown')}")

    gpu_info = env.get("gpu", {})
    if gpu_info.get("available"):
        for g in gpu_info.get("gpus", []):
            vram = f" ({g['vram_total_mb']}MB)" if g["vram_total_mb"] else ""
            parts.append(f"GPU: {g['name']}{vram}")
    else:
        parts.append("GPU: None detected")

    pkgs = env.get("installed_packages", {})
    if pkgs:
        pkg_str = ", ".join(f"{k}={v}" for k, v in sorted(pkgs.items()))
        parts.append(f"Packages: {pkg_str}")

    ros = env.get("ros_version")
    if ros:
        parts.append(f"ROS: {ros}")

    return "\n".join(parts)
