"""Troubleshooting knowledge base.

Pre-built decision trees for every common robotics training problem.
Each entry has symptom patterns, root cause, and specific fix commands.
"""

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class TroubleshootEntry:
    """A known problem with symptom, cause, and fix."""

    symptom: str
    keywords: list[str]  # for matching user queries
    cause: str
    fix: list[str]
    severity: str = "medium"  # low, medium, high


TROUBLESHOOTING_DB: list[TroubleshootEntry] = [
    TroubleshootEntry(
        symptom="GR00T training fails with 'EmbodimentTag not found'",
        keywords=["groot", "embodiment", "tag", "not found"],
        cause="Missing or incorrect embodiment tag in training command",
        fix=[
            "Use --embodiment-tag NEW_EMBODIMENT for custom robots",
            "Ensure modality.json exists in dataset's meta/ directory",
            "Generate with: orbit suggest <dataset> --framework groot",
        ],
        severity="high",
    ),
    TroubleshootEntry(
        symptom="GR00T says 'LeRobot v3 not supported'",
        keywords=["groot", "v3", "not supported", "version"],
        cause="GR00T N1.6 requires LeRobot v2, but data is v3",
        fix=[
            "Convert: orbit convert <dataset> --to lerobot-v2 -o ./groot-ready/",
            "Then add modality.json: orbit suggest <groot-ready> --framework groot",
        ],
        severity="high",
    ),
    TroubleshootEntry(
        symptom="OpenPI training fails with 'Config not found'",
        keywords=["openpi", "pi0", "config", "not found", "import"],
        cause="TrainConfig Python class not created or not importable",
        fix=[
            "Generate config: orbit suggest <dataset> --framework openpi",
            "Save to: openpi/configs/<name>_config.py",
            "Register in openpi/configs/__init__.py",
        ],
        severity="high",
    ),
    TroubleshootEntry(
        symptom="Loss is NaN after a few training steps",
        keywords=["nan", "loss", "diverge", "inf", "infinity", "explode"],
        cause="Usually data scaling issue or dead joints producing zero gradients",
        fix=[
            "Run: orbit analyze <dataset> to check for dead joints",
            "Run: orbit clean <dataset> to remove problematic episodes",
            "Try lower learning rate: divide current by 10",
            "Check action scaling — values should be roughly in [-1, 1]",
        ],
        severity="high",
    ),
    TroubleshootEntry(
        symptom="Loss plateaus and won't decrease further",
        keywords=["plateau", "stuck", "loss", "not decreasing", "flat", "stagnant"],
        cause="Conflicting demonstrations, multimodal action distributions, or data ceiling",
        fix=[
            "Run: orbit debug <training_dir> --ai -d <dataset>",
            "If high divergence: orbit curate <dataset> --budget <N>",
            "If low divergence but plateau: try Diffusion Policy instead of ACT",
            "If <50 episodes: collect more data (scaling laws say quantity matters)",
        ],
        severity="medium",
    ),
    TroubleshootEntry(
        symptom="Policy works in sim but fails on real robot",
        keywords=["sim", "real", "transfer", "gap", "sim2real", "works in simulation"],
        cause="Sim-to-real gap, usually visual or dynamics mismatch",
        fix=[
            "Check camera calibration matches training data",
            "Run: orbit analyze on real-robot data to check consistency",
            "Consider domain randomization or visual augmentation",
            "Verify action scaling matches between sim and real",
        ],
        severity="medium",
    ),
    TroubleshootEntry(
        symptom="Format conversion fails or produces empty dataset",
        keywords=["convert", "conversion", "empty", "failed", "format"],
        cause="Source format auto-detection failed or topics not found",
        fix=[
            "Specify format explicitly: orbit convert <src> --format <fmt>",
            "For rosbag: specify --action-topic and --state-topic",
            "For RLDS: check that .tfrecord files have correct schema",
            "Check: orbit analyze <src> --format <detected_format>",
        ],
        severity="medium",
    ),
    TroubleshootEntry(
        symptom="Training is extremely slow or OOM errors",
        keywords=["slow", "oom", "out of memory", "cuda", "memory", "killed"],
        cause="Batch size too large, image resolution too high, or insufficient VRAM",
        fix=[
            "Reduce batch_size by half",
            "Lower image resolution (e.g., 224x224 for ACT/Diffusion)",
            "Use gradient accumulation: double accumulation steps, halve batch size",
            "For OpenVLA: switch to LoRA (--lora_rank 32) to reduce VRAM",
            "For large models: use DeepSpeed ZeRO-2 or FSDP",
        ],
        severity="medium",
    ),
    TroubleshootEntry(
        symptom="Dataset loading is very slow",
        keywords=["loading", "slow", "data", "dataloader", "io", "bottleneck"],
        cause="Video decoding overhead, network I/O, or inefficient format",
        fix=[
            "Convert to LeRobot v3 (Parquet) for fastest local loading",
            "Use --num-workers 4 or higher in training",
            "If on HuggingFace Hub: download locally first with huggingface-cli",
            "For video datasets: pre-decode frames to numpy/parquet",
        ],
        severity="low",
    ),
    TroubleshootEntry(
        symptom="High action divergence across episodes",
        keywords=["divergence", "inconsistent", "diverge", "action", "variance"],
        cause="Demonstrations not consistent enough — different strategies for same task",
        fix=[
            "Run: orbit curate <dataset> --budget <N> to keep most consistent episodes",
            "Run: orbit clean <dataset> --aggressive to remove outliers",
            "Consider re-collecting data with stricter protocols",
            "Try Diffusion Policy which handles multimodal distributions better",
        ],
        severity="medium",
    ),
    TroubleshootEntry(
        symptom="Dead joints detected in dataset",
        keywords=["dead", "joint", "zero", "constant", "stuck"],
        cause="Some action dimensions never change — sensor fault or unused DOF",
        fix=[
            "Run: orbit analyze <dataset> to confirm which joints are dead",
            "If hardware fault: re-collect data",
            "If unused DOF: remove dimension from action space in config",
            "Run: orbit clean <dataset> to document and flag dead joints",
        ],
        severity="medium",
    ),
    TroubleshootEntry(
        symptom="Gradient norm is very high or spikes",
        keywords=["gradient", "spike", "explode", "norm", "high", "grad"],
        cause="Bad episodes causing large loss, or learning rate too high",
        fix=[
            "Enable gradient clipping: --gradient_clip_val 1.0",
            "Lower learning rate by 2-5x",
            "Run: orbit clean <dataset> to remove problematic episodes",
            "Check for outlier episodes with extreme actions",
        ],
        severity="medium",
    ),
    TroubleshootEntry(
        symptom="Checkpoint loading fails with size mismatch",
        keywords=["checkpoint", "size", "mismatch", "shape", "load", "state_dict"],
        cause="Model architecture changed between training and loading",
        fix=[
            "Verify action_dim and state_dim match between dataset and config",
            "Check if number of cameras changed (changes vision encoder input)",
            "For transfer: use --resume with compatible checkpoint only",
            "If fine-tuning: ensure base model matches checkpoint architecture",
        ],
        severity="high",
    ),
    TroubleshootEntry(
        symptom="Robot moves erratically or overshoots",
        keywords=["erratic", "overshoot", "jerk", "unstable", "robot", "movement"],
        cause="Action frequency mismatch, wrong action space, or poor data quality",
        fix=[
            "Check FPS: training data FPS must match deployment inference rate",
            "Verify action space: absolute vs delta, joint vs Cartesian",
            "Run: orbit analyze <dataset> to check temporal consistency",
            "Try chunk_size=1 for more reactive control",
        ],
        severity="medium",
    ),
    TroubleshootEntry(
        symptom="OSMO workflow fails to submit",
        keywords=["osmo", "workflow", "submit", "fail", "yaml"],
        cause="YAML schema errors, missing pool, or invalid platform",
        fix=[
            "Validate YAML: check indentation and required fields (name, image, command)",
            "Valid platforms: gb200, rtx-pro-6000, jetson-agx-thor",
            "Ensure --pool exists: osmo pool list",
            "Check image exists in container registry",
        ],
        severity="medium",
    ),
    TroubleshootEntry(
        symptom="OpenVLA training crashes with CUDA error",
        keywords=["openvla", "cuda", "crash", "error", "gpu"],
        cause="VRAM insufficient for full fine-tuning, or CUDA version mismatch",
        fix=[
            "Use LoRA: add --lora_rank 32 to reduce VRAM requirement to ~24GB",
            "Check CUDA version: OpenVLA OFT needs CUDA 11.8+",
            "Reduce batch_size to 4 or 2",
            "Try gradient checkpointing if available",
        ],
        severity="high",
    ),
]


def search_troubleshooting(query: str) -> list[TroubleshootEntry]:
    """Find relevant troubleshooting entries for a user query.

    Uses keyword matching — returns entries sorted by relevance (most keywords matched first).
    """
    query_lower = query.lower()
    scored: list[tuple[int, TroubleshootEntry]] = []

    for entry in TROUBLESHOOTING_DB:
        score = 0
        for kw in entry.keywords:
            if kw in query_lower:
                score += 1
        # Also check symptom text
        if any(word in entry.symptom.lower() for word in query_lower.split()):
            score += 1
        if score > 0:
            scored.append((score, entry))

    scored.sort(key=lambda x: x[0], reverse=True)
    return [entry for _, entry in scored]


def get_troubleshooting_summary() -> str:
    """Generate concise troubleshooting reference for the system prompt."""
    parts: list[str] = []
    for entry in TROUBLESHOOTING_DB:
        parts.append(f'### "{entry.symptom}"')
        parts.append(f"Cause: {entry.cause}")
        for fix in entry.fix:
            parts.append(f"  -> {fix}")
        parts.append("")
    return "\n".join(parts)
