"""Tool definitions and executors for ORBIT Assist.

Each tool maps to an ORBIT CLI command (or utility) executed via subprocess.
Tools are defined as Gemini-compatible function declarations and have matching
executor functions that run the actual commands.
"""

from __future__ import annotations

import json
import logging
import subprocess
import sys
from pathlib import Path

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Safety classification
# ---------------------------------------------------------------------------

SAFE_TOOLS = frozenset({
    "orbit_analyze",
    "orbit_gate",
    "orbit_suggest",
    "orbit_debug",
    "orbit_verify",
    "scan_directory",
    "search_orbit_docs",
})

DESTRUCTIVE_TOOLS = frozenset({
    "orbit_clean",
    "orbit_convert",
    "orbit_curate",
    "generate_config_file",
    "run_shell_command",
})

# ---------------------------------------------------------------------------
# Function declarations (Gemini function-calling format)
# ---------------------------------------------------------------------------

TOOL_DECLARATIONS: list[dict] = [
    {
        "name": "orbit_analyze",
        "description": (
            "Analyze a dataset's quality, consistency, and training readiness. "
            "Returns grade, score, episode count, action/state dimensions, "
            "consistency metrics, signal health, dead joints, scaling advice, "
            "and predicted success range. Always run this first on any dataset."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "dataset": {
                    "type": "string",
                    "description": "Path to dataset or HuggingFace repo ID (e.g. './data/' or 'lerobot/pusht')",
                },
                "format": {
                    "type": "string",
                    "enum": ["auto", "lerobot", "lerobot-local", "rlds", "hdf5", "rosbag", "directory"],
                    "description": "Force format detection (default: auto)",
                },
                "deep": {
                    "type": "boolean",
                    "description": "Run deep AI-powered analysis (slower but more thorough)",
                },
            },
            "required": ["dataset"],
        },
    },
    {
        "name": "orbit_curate",
        "description": (
            "Select the best N episodes from a dataset based on quality scoring. "
            "Returns selected/removed episodes, quality improvement, and budget curve."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "dataset": {
                    "type": "string",
                    "description": "Path to dataset",
                },
                "budget": {
                    "type": "integer",
                    "description": "Number of episodes to keep",
                },
                "method": {
                    "type": "string",
                    "enum": ["auto", "consistency", "mi"],
                    "description": "Scoring method (default: auto)",
                },
            },
            "required": ["dataset"],
        },
    },
    {
        "name": "orbit_gate",
        "description": (
            "CI/CD quality gate — pass/fail check for training readiness. "
            "Returns whether dataset passes, grade, score, predicted success range, "
            "and any blocking issues."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "dataset": {
                    "type": "string",
                    "description": "Path to dataset",
                },
                "policy": {
                    "type": "string",
                    "description": "Target policy (act, diffusion, smolvla, openvla, pi0, groot)",
                },
                "min_grade": {
                    "type": "string",
                    "enum": ["A", "B", "C", "D"],
                    "description": "Minimum acceptable grade (default: C)",
                },
            },
            "required": ["dataset", "policy"],
        },
    },
    {
        "name": "orbit_suggest",
        "description": (
            "Generate training commands for any policy framework. "
            "Returns ready-to-run commands with correct hyperparameters, "
            "conversion needs, and warnings."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "dataset": {
                    "type": "string",
                    "description": "Path to dataset",
                },
                "framework": {
                    "type": "string",
                    "enum": ["lerobot", "openvla", "openpi", "groot", "osmo", "all"],
                    "description": "Target framework (default: all)",
                },
                "policy": {
                    "type": "string",
                    "description": "Specific policy type (e.g. act, diffusion, pi0)",
                },
            },
            "required": ["dataset"],
        },
    },
    {
        "name": "orbit_convert",
        "description": (
            "Convert dataset between formats. Returns success status, "
            "output path, episode count, and any warnings."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "source": {
                    "type": "string",
                    "description": "Source dataset path",
                },
                "target_format": {
                    "type": "string",
                    "enum": ["lerobot-v2", "lerobot-v3", "hdf5"],
                    "description": "Target format",
                },
                "output": {
                    "type": "string",
                    "description": "Output directory path",
                },
            },
            "required": ["source", "target_format", "output"],
        },
    },
    {
        "name": "orbit_clean",
        "description": (
            "Remove bad episodes from a dataset. Returns removed/kept episode counts, "
            "removal reasons, and quality before/after."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "dataset": {
                    "type": "string",
                    "description": "Path to dataset",
                },
                "policy": {
                    "type": "string",
                    "description": "Target policy for cleaning decisions",
                },
                "aggressive": {
                    "type": "boolean",
                    "description": "Remove borderline episodes too (default: false)",
                },
            },
            "required": ["dataset"],
        },
    },
    {
        "name": "orbit_debug",
        "description": (
            "Diagnose training failures and trace them back to data issues. "
            "Returns training status, final loss, loss curve, root cause, "
            "data problems, and fix commands."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "training_dir": {
                    "type": "string",
                    "description": "Path to training output directory",
                },
                "dataset": {
                    "type": "string",
                    "description": "Optional dataset path for cross-referencing",
                },
                "ai": {
                    "type": "boolean",
                    "description": "Use AI-powered diagnosis (default: true)",
                },
            },
            "required": ["training_dir"],
        },
    },
    {
        "name": "orbit_verify",
        "description": (
            "Compare predicted quality against actual training outcome. "
            "Returns predicted grade/range, actual status, loss, prediction accuracy."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "training_dir": {
                    "type": "string",
                    "description": "Path to training output",
                },
                "dataset": {
                    "type": "string",
                    "description": "Dataset for prediction comparison",
                },
                "success_rate": {
                    "type": "number",
                    "description": "Manual success rate from rollout evaluation (0.0-1.0)",
                },
            },
            "required": ["training_dir"],
        },
    },
    {
        "name": "scan_directory",
        "description": (
            "Scan a directory for datasets, training outputs, configs, and environment info. "
            "Use this to discover what's in the current workspace."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "path": {
                    "type": "string",
                    "description": "Directory to scan (default: current directory)",
                },
            },
        },
    },
    {
        "name": "run_shell_command",
        "description": (
            "Execute a shell command. Use for pip installs, listing files, "
            "checking GPU status, running Python scripts, etc. "
            "Dangerous commands (rm -rf, sudo) are blocked."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "command": {
                    "type": "string",
                    "description": "Shell command to run",
                },
            },
            "required": ["command"],
        },
    },
    {
        "name": "generate_config_file",
        "description": (
            "Generate a framework-specific config file. "
            "For GR00T: modality.json and ModalityConfig. "
            "For OpenPI: TrainConfig Python class. "
            "For OSMO: workflow YAML."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "framework": {
                    "type": "string",
                    "enum": ["groot", "openpi", "osmo"],
                    "description": "Target framework",
                },
                "dataset_path": {
                    "type": "string",
                    "description": "Path to dataset (used to infer dimensions, cameras, etc.)",
                },
                "output_path": {
                    "type": "string",
                    "description": "Where to write the config file",
                },
            },
            "required": ["framework", "dataset_path", "output_path"],
        },
    },
    {
        "name": "search_orbit_docs",
        "description": (
            "Search ORBIT's internal knowledge base for answers about "
            "formats, frameworks, troubleshooting, and best practices."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Natural language question",
                },
            },
            "required": ["query"],
        },
    },
]


# ---------------------------------------------------------------------------
# Tool executors
# ---------------------------------------------------------------------------


def execute_tool(name: str, args: dict) -> dict:
    """Execute a tool by name with given arguments.

    Returns a dict with results (always JSON-serializable).
    """
    executor = _EXECUTORS.get(name)
    if not executor:
        return {"error": f"Unknown tool: {name}"}
    try:
        return executor(**args)
    except Exception as e:
        logger.debug("Tool %s failed: %s", name, e, exc_info=True)
        return {"error": f"{type(e).__name__}: {e}"}


def _resolve_python() -> str:
    """Resolve the correct Python executable for running orbit commands.

    Priority:
    1. The Python that's running THIS process (sys.executable) — this is
       almost always correct since orbit-assist IS an orbit command.
    2. If sys.executable doesn't have orbit installed (e.g. embedded Python),
       fall back to ``python3`` on PATH which inherits the user's active venv.
    """
    python = sys.executable
    # Sanity check: can this Python import orbit?
    # We skip the check if sys.executable is clearly a real Python.
    if not python or "freeze" in python.lower():
        python = "python3"
    return python


def _build_subprocess_env() -> dict[str, str]:
    """Build a clean subprocess environment.

    - Inherits the user's full environment (PATH, venv, CUDA, etc.)
    - Adds NO_COLOR=1 to suppress ANSI escape codes in output
    - Preserves VIRTUAL_ENV, CONDA_PREFIX, and PATH so the subprocess
      resolves the same packages as the parent process
    """
    import os

    env = os.environ.copy()
    env["NO_COLOR"] = "1"
    # Ensure PYTHONPATH doesn't interfere — the subprocess should use
    # the same site-packages as sys.executable
    # Don't remove PYTHONPATH entirely since the user may have set it
    # intentionally, but do ensure our Python can find orbit
    return env


def _run_orbit_command(args: list[str], timeout: int = 300) -> dict:
    """Run an orbit CLI command and capture output."""
    python = _resolve_python()
    cmd = [python, "-m", "orbit.cli"] + args
    logger.debug("Running: %s", " ".join(cmd))
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=timeout,
            env=_build_subprocess_env(),
        )
        stdout = result.stdout.strip()
        stderr = result.stderr.strip()

        # For gate command, exit code 1 means "fail" (not error) when --json
        if result.returncode != 0:
            # Try JSON parse — some commands use exit code 1 for "gate failed"
            parsed = _try_parse_json(stdout)
            if parsed is not None:
                parsed["exit_code"] = result.returncode
                return parsed
            return {
                "error": stderr or f"Command failed with exit code {result.returncode}",
                "stdout": stdout[:2000] if stdout else "",
                "exit_code": result.returncode,
            }

        # Try JSON parse
        parsed = _try_parse_json(stdout)
        if parsed is not None:
            return parsed
        return {"output": stdout[:4000]}
    except subprocess.TimeoutExpired:
        return {"error": f"Command timed out after {timeout}s"}


def _try_parse_json(text: str) -> dict | None:
    """Try to parse JSON from command output, handling ANSI prefixes/noise."""
    if not text:
        return None

    # First try: direct parse
    try:
        return json.loads(text)
    except (json.JSONDecodeError, ValueError):
        pass

    # Second try: find JSON object in the output (skip ANSI/log lines)
    # Look for the first { and last }
    start = text.find("{")
    end = text.rfind("}")
    if start >= 0 and end > start:
        try:
            return json.loads(text[start : end + 1])
        except (json.JSONDecodeError, ValueError):
            pass

    # Third try: find JSON array
    start = text.find("[")
    end = text.rfind("]")
    if start >= 0 and end > start:
        try:
            parsed = json.loads(text[start : end + 1])
            return {"data": parsed}
        except (json.JSONDecodeError, ValueError):
            pass

    return None


# --- Individual executors ---


def _exec_orbit_analyze(dataset: str, format: str = "auto", deep: bool = False) -> dict:
    cmd = ["analyze", dataset, "--json"]
    if format and format != "auto":
        cmd.extend(["--format", format])
    if deep:
        cmd.append("--deep")
    return _run_orbit_command(cmd, timeout=300)


def _exec_orbit_curate(dataset: str, budget: int | None = None, method: str = "auto") -> dict:
    cmd = ["curate", dataset, "--json"]
    if budget:
        cmd.extend(["--budget", str(budget)])
    if method and method != "auto":
        cmd.extend(["--method", method])
    return _run_orbit_command(cmd, timeout=300)


def _exec_orbit_gate(dataset: str, policy: str, min_grade: str = "C") -> dict:
    cmd = ["gate", dataset, "-p", policy, "--min-grade", min_grade, "--json"]
    return _run_orbit_command(cmd, timeout=300)


def _exec_orbit_suggest(
    dataset: str, framework: str = "all", policy: str | None = None
) -> dict:
    cmd = ["suggest", dataset, "--json"]
    if framework and framework != "all":
        cmd.extend(["--framework", framework])
    if policy:
        cmd.extend(["--policy", policy])
    return _run_orbit_command(cmd, timeout=120)


def _exec_orbit_convert(source: str, target_format: str, output: str) -> dict:
    cmd = ["convert", source, "--to", target_format, "-o", output]
    return _run_orbit_command(cmd, timeout=600)


def _exec_orbit_clean(
    dataset: str, policy: str = "act", aggressive: bool = False
) -> dict:
    cmd = ["clean", dataset, "-p", policy, "--json"]
    if aggressive:
        cmd.append("--aggressive")
    return _run_orbit_command(cmd, timeout=300)


def _exec_orbit_debug(
    training_dir: str, dataset: str | None = None, ai: bool = True
) -> dict:
    # When using --json, skip --ai since AI diagnosis prints to console, not JSON.
    # The JSON output contains rule-based findings which are sufficient for the agent.
    cmd = ["debug", training_dir, "--json"]
    if dataset:
        cmd.extend(["-d", dataset])
    return _run_orbit_command(cmd, timeout=120)


def _exec_orbit_verify(
    training_dir: str, dataset: str | None = None, success_rate: float | None = None
) -> dict:
    cmd = ["verify", training_dir, "--json"]
    if dataset:
        cmd.extend(["-d", dataset])
    if success_rate is not None:
        cmd.extend(["--success-rate", str(success_rate)])
    return _run_orbit_command(cmd, timeout=120)


def _exec_scan_directory(path: str = ".") -> dict:
    from orbit.assist.context import build_workspace_context

    return build_workspace_context(path)


def _exec_run_shell_command(command: str) -> dict:
    """Execute a shell command with safety checks."""
    # Block dangerous commands
    dangerous = ["rm -rf /", "rm -rf ~", "sudo rm", "mkfs", "dd if=", "> /dev/", "chmod 777 /"]
    cmd_lower = command.lower().strip()
    for pattern in dangerous:
        if pattern in cmd_lower:
            return {"error": f"Blocked dangerous command pattern: {pattern}"}

    try:
        result = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=60,
        )
        return {
            "exit_code": result.returncode,
            "stdout": result.stdout[:4000],
            "stderr": result.stderr[:2000] if result.stderr else "",
        }
    except subprocess.TimeoutExpired:
        return {"error": "Command timed out after 60s"}


def _exec_generate_config_file(
    framework: str, dataset_path: str, output_path: str
) -> dict:
    """Generate a framework-specific config file."""
    # First, get dataset info via orbit analyze
    info = _exec_orbit_analyze(dataset_path)
    if "error" in info:
        return {"error": f"Could not analyze dataset: {info['error']}"}

    if framework == "groot":
        return _generate_groot_config(info, dataset_path, output_path)
    elif framework == "openpi":
        return _generate_openpi_config(info, dataset_path, output_path)
    elif framework == "osmo":
        return _generate_osmo_workflow(info, dataset_path, output_path)
    else:
        return {"error": f"Config generation not supported for framework: {framework}"}


def _generate_groot_config(info: dict, dataset_path: str, output_path: str) -> dict:
    """Generate GR00T modality.json and ModalityConfig."""
    # Extract dimensions from analysis
    shapes = info.get("shapes", {})
    cameras = info.get("cameras", [])

    modality = {
        "modality": {
            "action": {"dim": info.get("action_dim", 7)},
            "state": {"dim": info.get("state_dim", 7)},
        },
        "cameras": cameras if cameras else ["front"],
    }

    # Write modality.json to dataset meta dir
    meta_dir = Path(dataset_path) / "meta"
    modality_json_path = meta_dir / "modality.json"
    try:
        meta_dir.mkdir(parents=True, exist_ok=True)
        modality_json_path.write_text(json.dumps(modality, indent=2))
    except OSError as e:
        return {"error": f"Failed to write modality.json: {e}"}

    # Generate ModalityConfig Python file
    config_content = (
        "from gr00t.experiment.data_config import ModalityConfig\n\n"
        f"MODALITY_CONFIG = ModalityConfig(\n"
        f'    action_dim={info.get("action_dim", 7)},\n'
        f'    state_dim={info.get("state_dim", 7)},\n'
        f"    cameras={cameras if cameras else ['front']},\n"
        f")\n"
    )

    try:
        Path(output_path).write_text(config_content)
    except OSError as e:
        return {"error": f"Failed to write config: {e}"}

    return {
        "modality_json": str(modality_json_path),
        "modality_config": output_path,
        "content_preview": config_content,
        "instructions": (
            f"1. modality.json written to {modality_json_path}\n"
            f"2. ModalityConfig written to {output_path}\n"
            f"3. Use --modality-config-path {output_path} in your GR00T training command"
        ),
    }


def _generate_openpi_config(info: dict, dataset_path: str, output_path: str) -> dict:
    """Generate OpenPI TrainConfig."""
    config_name = Path(output_path).stem
    config_content = (
        "from openpi.training.config import TrainConfig\n"
        "from openpi.models.pi0 import Pi0Config\n\n\n"
        f"def {config_name}() -> TrainConfig:\n"
        f"    return TrainConfig(\n"
        f"        model=Pi0Config(),\n"
        f'        dataset="{dataset_path}",\n'
        f"        num_train_steps=50000,\n"
        f"        batch_size=32,\n"
        f"    )\n"
    )

    try:
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        Path(output_path).write_text(config_content)
    except OSError as e:
        return {"error": f"Failed to write config: {e}"}

    return {
        "file_path": output_path,
        "content_preview": config_content,
        "instructions": (
            f"1. Config written to {output_path}\n"
            f"2. Register in openpi/configs/__init__.py\n"
            f"3. Train with: XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py {config_name}"
        ),
    }


def _generate_osmo_workflow(info: dict, dataset_path: str, output_path: str) -> dict:
    """Generate OSMO workflow YAML."""
    import textwrap

    workflow_content = textwrap.dedent(f"""\
        workflow:
          tasks:
          - name: quality-gate
            image: orbit-robotics:latest
            command: orbit gate {dataset_path} -p act --min-grade B

          - name: train-policy
            image: nvcr.io/nvidia/pytorch:latest
            platform: rtx-pro-6000
            resources:
              gpu: 1
            inputs:
            - task: quality-gate
            command: |
              pip install orbit-robotics[formats]
              # Replace with your actual training command
              python train.py --dataset {dataset_path}

          - name: verify-result
            image: orbit-robotics:latest
            command: orbit verify /output/ -d {dataset_path}
            inputs:
            - task: train-policy
    """)

    try:
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        Path(output_path).write_text(workflow_content)
    except OSError as e:
        return {"error": f"Failed to write workflow: {e}"}

    return {
        "file_path": output_path,
        "content_preview": workflow_content,
        "instructions": (
            f"1. Workflow written to {output_path}\n"
            f"2. Edit the train-policy command for your specific setup\n"
            f"3. Submit with: osmo workflow submit {output_path} --pool <your-pool>"
        ),
    }


def _exec_search_orbit_docs(query: str) -> dict:
    """Search the internal knowledge base."""
    from orbit.assist.knowledge.troubleshoot import search_troubleshooting

    results = search_troubleshooting(query)

    entries = []
    for entry in results[:5]:
        entries.append({
            "symptom": entry.symptom,
            "cause": entry.cause,
            "fix": entry.fix,
        })

    # Also check framework knowledge
    from orbit.assist.knowledge.frameworks import FRAMEWORK_KNOWLEDGE

    related_commands: list[str] = []
    query_lower = query.lower()
    for fw_key, fw in FRAMEWORK_KNOWLEDGE.items():
        if fw_key in query_lower or fw["name"].lower() in query_lower:
            related_commands.append(f"orbit suggest <dataset> --framework {fw_key}")

    if "analyze" in query_lower or "quality" in query_lower:
        related_commands.append("orbit analyze <dataset> --json")
    if "convert" in query_lower or "format" in query_lower:
        related_commands.append("orbit convert <source> --to <format> -o <output>")
    if "debug" in query_lower or "fail" in query_lower or "error" in query_lower:
        related_commands.append("orbit debug <training_dir> --ai")
    if "clean" in query_lower:
        related_commands.append("orbit clean <dataset>")
    if "curate" in query_lower or "budget" in query_lower:
        related_commands.append("orbit curate <dataset> --budget <N>")

    return {
        "relevant_entries": entries,
        "related_commands": related_commands,
    }


# ---------------------------------------------------------------------------
# Executor registry
# ---------------------------------------------------------------------------

_EXECUTORS: dict[str, callable] = {
    "orbit_analyze": _exec_orbit_analyze,
    "orbit_curate": _exec_orbit_curate,
    "orbit_gate": _exec_orbit_gate,
    "orbit_suggest": _exec_orbit_suggest,
    "orbit_convert": _exec_orbit_convert,
    "orbit_clean": _exec_orbit_clean,
    "orbit_debug": _exec_orbit_debug,
    "orbit_verify": _exec_orbit_verify,
    "scan_directory": _exec_scan_directory,
    "run_shell_command": _exec_run_shell_command,
    "generate_config_file": _exec_generate_config_file,
    "search_orbit_docs": _exec_search_orbit_docs,
}
