from __future__ import annotations

import json
import os
import sys
import time
import threading
import socket
import platform
from typing import Any, Dict, Optional
from pathlib import Path

import click

from actor.cli.utils import (
    CLIContext,
    decode_connection_token,
    generate_config_template,
    get_log_file_path,
    is_sensitive_key,
    load_config,
    redact,
    resolve_redis_password,
    sanitize_error_message,
)
from actor import __version__
from actor.identity import load_or_create_instance_id
from actor.logging import LoggingService
from actor.queue import QueueService, QueueError
from actor.runtime import RuntimeConfig
from actor.runtime.task_router import TaskRouter
from actor.observability import ObservabilityService
from actor.authentication.exceptions import RedisConnectionError


def _build_logging(log_dir: Optional[str], level: str, log_format: str) -> LoggingService:
    import logging

    level_map = {"DEBUG": logging.DEBUG, "INFO": logging.INFO, "WARNING": logging.WARNING, "ERROR": logging.ERROR}
    svc = LoggingService(log_dir=log_dir, level=level_map.get(level, logging.INFO), log_format=log_format)
    return svc


def _get_config_cache_path(agent_key: str) -> Path:
    # Default to ~/.causum/agent/<agent_key>/config_cache.json
    base = Path.home() / ".causum" / "agent" / agent_key
    base.mkdir(parents=True, exist_ok=True)
    try:
        os.chmod(base, 0o700)
    except OSError:
        pass
    return base / "config_cache.json"


def _load_cached_config(cache_path: Path, logger) -> Optional[Dict[str, Any]]:
    if not cache_path.exists():
        return None
    try:
        data = json.loads(cache_path.read_text())
        if isinstance(data, dict):
            logger.info("Loaded cached config from %s", cache_path)
            return data
    except Exception as exc:
        logger.warning("Failed to load cached config from %s: %s", cache_path, exc)
    return None


def _sanitize_config_for_cache(cfg_msg: Dict[str, Any]) -> Dict[str, Any]:
    """Strip raw secrets before caching; keep secret acquisition methods."""
    scrub_fields = {
        "api_key",
        "password",
        "token",
        "access_key_id",
        "secret_access_key",
        "private_key",
        "client_secret",
        "jwt_token",
    }
    cfg = json.loads(json.dumps(cfg_msg))
    for section in ("database", "llm"):
        sec = cfg.get(section)
        if not isinstance(sec, dict):
            continue
        for key in list(sec.keys()):
            if key in scrub_fields:
                sec.pop(key, None)
        # Drop database.url if it likely embeds credentials
        if section == "database" and isinstance(sec.get("url"), str):
            try:
                from urllib.parse import urlparse

                parsed = urlparse(sec["url"])
                if parsed.username or parsed.password:
                    sec.pop("url", None)
            except Exception:
                pass
    return cfg


def _save_cached_config(cache_path: Path, cfg_msg: Dict[str, Any], logger) -> None:
    try:
        cache_path.write_text(json.dumps(_sanitize_config_for_cache(cfg_msg)))
        try:
            os.chmod(cache_path, 0o600)
        except OSError:
            pass
        logger.info("Cached config saved to %s", cache_path)
    except Exception as exc:
        logger.warning("Failed to save cached config to %s: %s", cache_path, exc)


def _build_redis_client(redis_url: str, password: Optional[str]):
    try:
        import redis
    except ImportError as exc:  # pragma: no cover - external dep
        raise RuntimeError("redis library is required for queue operations") from exc
    return redis.Redis.from_url(redis_url, password=password, decode_responses=True)


def _ensure_db_driver(db_url: str, logger) -> None:
    """Fail fast if required SQLAlchemy driver is missing."""
    try:
        from sqlalchemy.engine import make_url
    except Exception as exc:  # pragma: no cover
        logger.error("SQLAlchemy required to parse DB URL: %s", exc)
        raise
    url = make_url(db_url)
    scheme = url.drivername.lower()
    driver_required = None
    extra_hint = None
    if scheme.startswith("postgres"):
        driver_required = "psycopg2"
        extra_hint = "db-postgres"
    elif scheme.startswith("mysql") or scheme.startswith("mariadb"):
        driver_required = "pymysql"
        extra_hint = "db-mysql"
    elif scheme.startswith("mssql"):
        driver_required = "pyodbc"
        extra_hint = "db-mssql"
    elif scheme.startswith("snowflake"):
        driver_required = "snowflake.sqlalchemy"
        extra_hint = "db-snowflake"
    elif scheme.startswith("oracle"):
        driver_required = "oracledb"
        extra_hint = "db-oracle"
    elif scheme.startswith("trino"):
        driver_required = "trino"
        extra_hint = "db-trino"
    elif scheme.startswith("bigquery"):
        driver_required = "pybigquery"
        extra_hint = "db-bigquery"
    elif scheme.startswith("awsathena"):
        driver_required = "pyathena"
        extra_hint = "db-athena"
    elif scheme.startswith("duckdb"):
        driver_required = "duckdb_engine"
        extra_hint = "db-duckdb"
    elif scheme.startswith("databricks"):
        driver_required = "databricks.sql"
        extra_hint = "db-databricks"
    elif scheme.startswith("sqlite"):
        driver_required = None  # builtin
    if driver_required:
        try:
            __import__(driver_required)
        except ImportError as exc:
            hint = f" (install with `pip install .[{extra_hint}]`)" if extra_hint else ""
            raise RuntimeError(f"Database driver '{driver_required}' is required for URL '{db_url}'{hint}") from exc


def _validate_connections(config: Dict[str, Any]) -> Dict[str, Any]:
    results: Dict[str, Any] = {}
    redis_cfg = config.get("redis", {})
    redis_url = redis_cfg.get("url")
    password = resolve_redis_password(redis_cfg)
    if redis_url:
        try:
            client = _build_redis_client(redis_url, password)
            latency = client.ping()
            results["redis"] = {"status": "healthy", "ping": latency}
        except Exception as exc:
            results["redis"] = {"status": "unhealthy", "error": str(exc)}
    else:
        results["redis"] = {"status": "unknown", "error": "redis url missing"}
    return results


def _apply_config_and_validate(cfg_msg: Dict[str, Any], runtime_cfg: RuntimeConfig, logger, obs: ObservabilityService) -> None:
    """Parse, schema-validate, business-validate, and apply config.

    Connector validation is intentionally deferred to the caller (the ``start``
    command validates via the real TaskRouter after config is applied).
    """
    from actor.runtime.config_store import validate_config_payload, validate_config_schema

    if not isinstance(cfg_msg, dict):
        logger.error("Config message is not a dict (type=%s): %r", type(cfg_msg).__name__, cfg_msg)
        raise ValueError("Config message must be an object")
    new_version = cfg_msg.get("version")
    current_version = getattr(runtime_cfg, "version", None)
    try:
        new_version_cmp = int(new_version) if new_version is not None else None
    except (TypeError, ValueError):
        new_version_cmp = None
    if new_version_cmp is not None and current_version is not None and new_version_cmp == current_version:
        logger.info("Config version unchanged (%s); skipping apply", new_version)
        return
    validate_config_schema(cfg_msg)
    validate_config_payload(cfg_msg)
    runtime_cfg.update_from_config_message(cfg_msg)
    logger.info("Applied config to runtime (version=%s)", new_version if new_version is not None else runtime_cfg.version)


@click.group()
def cli():
    """Agent CLI - On-prem worker for your SaaS."""
    pass


@cli.command()
@click.option("--config", type=click.Path(exists=True), help="Path to config file")
@click.option("--log-level", type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR"]), default="INFO")
@click.option("--log-format", type=click.Choice(["json", "text"]), default="json")
@click.option("--log-dir", type=click.Path(), help="Directory to write logs (optional)")
@click.option("--no-validation", is_flag=True, help="Skip startup validation")
@click.option("--instance-name", help="Optional operator-friendly name (display only)")
def start(config, log_level, log_format, log_dir, no_validation, instance_name):
    """Start the agent worker (minimal stub until connectors wired)."""
    try:
        cfg = load_config(config_path=config)
    except Exception as exc:
        click.echo(f"Configuration error: {exc}", err=True)
        sys.exit(1)

    log_dir_resolved = log_dir or (cfg.get("logging") or {}).get("log_dir")
    log_service = _build_logging(log_dir=log_dir_resolved, level=log_level, log_format=log_format)
    logger = log_service.get_logger("agent.start")
    logger.info("Starting agent CLI stub")

    if not no_validation:
        v = _validate_connections(cfg)
        logger.info("Validation results: %s", v)

    redis_cfg = cfg.get("redis", {})
    agent_cfg = cfg.get("agent", {})
    if config:
        logger.info("Loaded config from path: %s", config)
    elif os.getenv("AGENT_KEY"):
        logger.info("Loaded config from AGENT_KEY")
    elif os.getenv("AGENT_CONNECTION_TOKEN"):
        logger.info("Loaded config from AGENT_CONNECTION_TOKEN")
    else:
        logger.warning("Config source unknown (no path or env token detected)")
    redis_url = redis_cfg.get("url")
    password = resolve_redis_password(redis_cfg)

    if not redis_url or not agent_cfg.get("key"):
        logger.error("Missing redis url or agent key; cannot start queue consumer.")
        sys.exit(1)

    instance_name = instance_name or agent_cfg.get("instance_name")
    configured_instance_id = agent_cfg.get("instance_id")
    instance_id, identity_path = load_or_create_instance_id(
        log_dir=log_dir_resolved,
        configured_id=configured_instance_id,
        agent_key=agent_cfg.get("key"),
    )
    logger.info("Using instance_id=%s (identity_file=%s)", instance_id, identity_path)
    logger.info("Redis target: %s", redis_url)
    logger.info("Agent key: %s", agent_cfg.get("key"))

    try:
        redis_client = _build_redis_client(redis_url, password)
        try:
            redis_client.ping()
            logger.info("Redis connection OK")
        except Exception as exc:
            raise RedisConnectionError(f"Redis ping failed: {exc}") from exc
        queue = QueueService(
            agent_key=agent_cfg["key"],
            instance_id=instance_id,
            redis_client=redis_client,
            redis_url=redis_url,
            redis_password=password,
        )
        queue.initialize_consumer()
        cache_path = _get_config_cache_path(agent_cfg["key"])
        try:
            registration_metadata = {
                "hostname": socket.gethostname(),
                "platform": platform.system().lower(),
                "python_version": platform.python_version(),
            }
            # Capture cursor before publishing so we don't miss a fast ack
            results_cursor = queue.capture_results_cursor()
            queue.publish_instance_registration(instance_name=instance_name, version=__version__, metadata=registration_metadata)
            reg_ack = queue.wait_for_registration_ack(instance_id=instance_id, timeout=30, start_id=results_cursor)
        except Exception as exc:
            logger.error("Failed to publish instance registration: %s", exc)
            sys.exit(1)
        # If the ack was actually a config message, apply it immediately to avoid races.
        initial_cfg_msg = None
        if isinstance(reg_ack, dict) and reg_ack.get("type") == "config":
            initial_cfg_msg = QueueService._parse_config_payload(reg_ack)

        stop_event = threading.Event()

        results_stream = f"agent:{agent_cfg['key']}:results"
        telemetry_stream = f"agent:{agent_cfg['key']}:telemetry"
        obs = ObservabilityService(
            instance_id=instance_id,
            instance_name=instance_name,
            version=__version__,
            results_stream=results_stream,
            telemetry_stream=telemetry_stream,
            redis_connection=queue._conn,
        )

        runtime_cfg = RuntimeConfig()
        config_ready = False
        if initial_cfg_msg:
            try:
                _apply_config_and_validate(initial_cfg_msg, runtime_cfg, logger, obs)
                obs.set_config_version(runtime_cfg.get_version())
                obs.publish_config_applied(runtime_cfg.get_version())
                logger.info("Received and validated initial config from SaaS (via registration ack)")
                _save_cached_config(cache_path, initial_cfg_msg, logger)
                config_ready = True
            except Exception as exc:
                logger.error("Failed to apply config from registration ack: %s", exc)
        if not config_ready:
            cached = _load_cached_config(cache_path, logger)
            if cached:
                try:
                    _apply_config_and_validate(cached, runtime_cfg, logger, obs)
                    obs.set_config_version(runtime_cfg.get_version())
                    logger.info("Applied cached config; awaiting SaaS refresh")
                    config_ready = True
                except Exception as exc:
                    logger.warning("Cached config invalid; ignoring: %s", exc)
        if not config_ready:
            try:
                queue.publish_config_request()
                cfg_msg = queue.wait_for_config(timeout=30)
                _apply_config_and_validate(cfg_msg, runtime_cfg, logger, obs)
                obs.set_config_version(runtime_cfg.get_version())
                obs.publish_config_applied(runtime_cfg.get_version())
                logger.info("Received and validated initial config from SaaS")
                _save_cached_config(cache_path, cfg_msg, logger)
                config_ready = True
            except Exception as exc:
                logger.warning("No config received from SaaS; continuing without config: %s", exc)

        # DB driver preflight if URL provided
        if config_ready:
            db_cfg = runtime_cfg.get_database()
            if db_cfg and db_cfg.url:
                try:
                    _ensure_db_driver(db_cfg.url, logger)
                except Exception as exc:
                    logger.error("Database driver missing: %s", exc)
                    sys.exit(1)

        router = TaskRouter(runtime_config=runtime_cfg, observability=obs)

        def _run_entitlement_validation() -> None:
            if runtime_cfg.get_database():
                try:
                    router._ensure_db().validate()
                except Exception as exc:
                    obs.publish_entitlement_error("database", sanitize_error_message(exc))
                    raise
            if runtime_cfg.get_llm():
                try:
                    llm = router._ensure_llm()
                    if llm:
                        llm.validate()
                except Exception as exc:
                    obs.publish_entitlement_error("llm", sanitize_error_message(exc))
                    raise
            entitlements_ok = []
            if runtime_cfg.get_database():
                entitlements_ok.append("database")
            if runtime_cfg.get_llm():
                entitlements_ok.append("llm")
            obs.publish_entitlements_ok(entitlements_ok)
            obs.send_validation_report()

        # Run validation after config
        if config_ready:
            try:
                _run_entitlement_validation()
            except Exception as exc:
                logger.error("Validation flow failed: %s", exc)
        else:
            logger.info("Waiting for config update; tasks may be rejected until config arrives.")

        logger.info("Agent ready - consuming tasks")

        # Start heartbeat/metrics loops
        obs.start_heartbeat(interval=30)

        def config_watch_loop():
            backoff = 2
            while not stop_event.is_set():
                try:
                    updates = queue.poll_config_updates(block_ms=5000, count=5)
                    for cfg_msg in updates:
                        new_version = cfg_msg.get("version")
                        try:
                            new_version_int = int(new_version) if new_version is not None else None
                        except (TypeError, ValueError):
                            new_version_int = None
                        if new_version_int is not None and new_version_int == runtime_cfg.get_version():
                            logger.debug("Config version %s unchanged; skipping", new_version)
                            continue
                        logger.info("Config update received; applying")
                        _apply_config_and_validate(cfg_msg, runtime_cfg, logger, obs)
                        obs.set_config_version(runtime_cfg.get_version())
                        try:
                            _run_entitlement_validation()
                        except Exception as exc:
                            logger.error("Validation flow failed after config update: %s", exc)
                        obs.publish_config_applied(runtime_cfg.get_version())
                        _save_cached_config(cache_path, cfg_msg, logger)
                    backoff = 2
                except Exception as exc:
                    logger.error("Config watch failed: %s", exc)
                    stop_event.wait(backoff)
                    backoff = min(backoff * 2, 60)

        cfg_thread = threading.Thread(target=config_watch_loop, daemon=True)
        cfg_thread.start()

        def backlog_loop():
            while not stop_event.is_set():
                try:
                    queue.publish_backlog_report()
                except Exception as exc:
                    logger.error("Backlog report failed: %s", exc)
                stop_event.wait(300)

        backlog_thread = threading.Thread(target=backlog_loop, daemon=True)
        backlog_thread.start()

        shutdown_event = threading.Event()
        shutdown_cleanup_error = threading.Event()

        def _handle(task: Dict[str, Any]) -> Dict[str, Any]:
            try:
                if task.get("task_type") == "shutdown":
                    logger.info("Shutdown task received; stopping agent.")
                    try:
                        router.shutdown()
                    except Exception as exc:
                        logger.error("Shutdown cleanup failed: %s", exc)
                        shutdown_cleanup_error.set()
                    # Remove cached config + identity to force fresh bootstrap next start.
                    try:
                        if cache_path.exists():
                            cache_path.unlink()
                            logger.info("Deleted config cache at %s", cache_path)
                    except Exception as exc:
                        logger.error("Failed to delete config cache: %s", exc)
                        shutdown_cleanup_error.set()
                    try:
                        if identity_path.exists():
                            identity_path.unlink()
                            logger.info("Deleted identity file at %s", identity_path)
                    except Exception as exc:
                        logger.error("Failed to delete identity file: %s", exc)
                        shutdown_cleanup_error.set()
                    stop_event.set()
                    obs.stop_heartbeat()
                    queue.stop()
                    shutdown_event.set()
                    return {"status": "ok" if not shutdown_cleanup_error.is_set() else "error"}
                return router.handle(task)
            except Exception as exc:
                logger.error("Task handling failed: %s", exc)
                return {"status": "error", "error": sanitize_error_message(exc)}

        try:
            queue.consume_tasks(callback=_handle, block_ms=5000)
        except KeyboardInterrupt:
            logger.info("Shutdown requested, stopping consumer.")
            queue.stop()
        finally:
            stop_event.set()
            obs.stop_heartbeat()
            if shutdown_event.is_set():
                logger.info("Agent shutdown task completed; exiting.")
    except RedisConnectionError as exc:
        logger.error("Redis connection failed: %s", exc)
        sys.exit(2)
    except Exception as exc:
        logger.error("Runtime error: %s", exc)
        sys.exit(3)


@cli.command()
@click.option("--component", type=click.Choice(["redis", "database", "llm", "pipeline"]), help="Component to validate")
@click.option("--format", "out_format", type=click.Choice(["text", "json"]), default="text")
def validate(component, out_format):
    """Validate agent configuration and connections."""
    try:
        cfg = load_config()
        redis_cfg = cfg.get("redis", {})
        agent_cfg = cfg.get("agent", {})
        redis_url = redis_cfg.get("url")
        password = resolve_redis_password(redis_cfg)
        if not redis_url or not agent_cfg.get("key"):
            raise RuntimeError("redis url/agent key required")
        instance_name = agent_cfg.get("instance_name")
        configured_instance_id = agent_cfg.get("instance_id")
        instance_id, _ = load_or_create_instance_id(
            log_dir=(cfg.get("logging") or {}).get("log_dir"),
            configured_id=configured_instance_id,
            agent_key=agent_cfg.get("key"),
        )
        redis_client = _build_redis_client(redis_url, password)
        try:
            redis_client.ping()
        except Exception as exc:
            raise RedisConnectionError(f"Redis ping failed: {exc}") from exc
        queue = QueueService(agent_key=agent_cfg["key"], instance_id=instance_id, redis_client=redis_client, redis_url=redis_url, redis_password=password)
        queue.initialize_consumer()
        queue.publish_config_request()
        cfg_msg = queue.wait_for_config(timeout=10)
        from actor.runtime.config_store import validate_config_payload

        validate_config_payload(cfg_msg)
        runtime_cfg = RuntimeConfig()
        runtime_cfg.update_from_config_message(cfg_msg)
        results = {}
        # Redis ping
        results["redis"] = {"status": "healthy", "ping": True}
        obs = ObservabilityService(
            instance_id=instance_id,
            instance_name=instance_name,
            version=__version__,
            redis_client=redis_client,
            results_stream=f"agent:{agent_cfg['key']}:results",
            telemetry_stream=f"agent:{agent_cfg['key']}:telemetry",
        )
        router = TaskRouter(runtime_config=runtime_cfg, observability=obs)
        if not component or component == "database":
            if runtime_cfg.get_database():
                try:
                    router._ensure_db().validate()
                    results["database"] = {"status": "healthy"}
                except Exception as exc:
                    results["database"] = {"status": "unhealthy", "error": str(exc)}
        if not component or component == "llm":
            if runtime_cfg.get_llm():
                try:
                    llm = router._ensure_llm()
                    if llm:
                        llm.validate()
                    results["llm"] = {"status": "healthy"}
                except Exception as exc:
                    results["llm"] = {"status": "unhealthy", "error": str(exc)}
        if component == "pipeline":
            try:
                pipeline_result = router._handle_pipeline_validate()
                results["pipeline"] = pipeline_result
            except Exception as exc:
                results["pipeline"] = {"status": "unhealthy", "error": str(exc)}

        if out_format == "json":
            click.echo(json.dumps(results, indent=2))
        else:
            for name, res in results.items():
                status = res.get("status")
                msg = res.get("error") or res
                click.echo(f"{name}: {status} ({msg})")
        status_values = {r.get("status") for r in results.values()}
        if "unhealthy" in status_values:
            sys.exit(2)
    except Exception as exc:
        click.echo(f"Validation failed: {exc}", err=True)
        sys.exit(1)


@cli.group()
def config():
    """Configuration management commands."""
    pass


@config.command("show")
@click.option("--format", "out_format", type=click.Choice(["text", "json"]), default="text")
@click.option("--show-secrets", is_flag=True, help="Show unredacted secrets (use carefully)")
def config_show(out_format, show_secrets):
    try:
        cfg = load_config()
        def redact_value(v):
            return v if show_secrets else redact(str(v))
        if out_format == "json":
            redacted = json.loads(json.dumps(cfg))
            if not show_secrets:
                for section in ("redis", "agent"):
                    if section in redacted and isinstance(redacted[section], dict):
                        for k in list(redacted[section].keys()):
                            if is_sensitive_key(k):
                                redacted[section][k] = redact_value(redacted[section][k])
            click.echo(json.dumps(redacted, indent=2))
        else:
            for section, values in cfg.items():
                click.echo(f"{section}:")
                if isinstance(values, dict):
                    for k, v in values.items():
                        if not show_secrets and is_sensitive_key(k):
                            v = redact_value(v)
                        click.echo(f"  {k}: {v}")
    except Exception as exc:
        click.echo(f"Error: {exc}", err=True)
        sys.exit(1)


@config.command("dump")
@click.option("--format", "out_format", type=click.Choice(["yaml", "json"]), default="yaml")
@click.option("--show-secrets", is_flag=True, help="Show unredacted secrets (use carefully)")
def config_dump(out_format, show_secrets):
    """Dump the current config for operator customization."""
    try:
        cfg = load_config()
        if not show_secrets:
            def _redact_dict(d):
                for k, v in d.items():
                    if isinstance(v, dict):
                        _redact_dict(v)
                    elif isinstance(v, str) and is_sensitive_key(k):
                        d[k] = redact(v)
            _redact_dict(cfg)
        if out_format == "json":
            click.echo(json.dumps(cfg, indent=2))
        else:
            import yaml
            click.echo(yaml.safe_dump(cfg, sort_keys=False))
    except Exception as exc:
        click.echo(f"Error: {exc}", err=True)
        sys.exit(1)


@config.command("test")
@click.option("--token", help="Connection token to test")
@click.option("--show-decoded", is_flag=True)
def config_test(token, show_decoded):
    try:
        token = token or os.getenv("AGENT_KEY") or os.getenv("AGENT_CONNECTION_TOKEN")
        if not token:
            click.echo("No token provided", err=True)
            sys.exit(1)
        decoded = decode_connection_token(token)
        click.echo("✓ Token format: Valid")
        click.echo("✓ Token decoding: Success")
        if show_decoded:
            redacted = {k: (redact(str(v)) if is_sensitive_key(k) else v) for k, v in decoded.items()}
            click.echo(json.dumps(redacted, indent=2))
    except Exception as exc:
        click.echo(f"Token test failed: {exc}", err=True)
        sys.exit(1)


@config.command("generate")
@click.option("--with-examples", is_flag=True)
def config_generate(with_examples):
    click.echo(generate_config_template(with_examples=with_examples))


@cli.command()
@click.option("--format", "out_format", type=click.Choice(["text", "json"]), default="text")
@click.option("--watch", is_flag=True, help="Refresh every 5 seconds")
def health(out_format, watch):
    """Check agent health status (uses Redis streams if available)."""
    try:
        ctx = CLIContext()
        ctx.require_redis_and_key()
        client = ctx.build_redis_client()
        results_stream = f"agent:{ctx.agent_key}:results"
        telemetry_stream = f"agent:{ctx.agent_key}:telemetry"
        def fetch():
            msgs = client.xrevrange(telemetry_stream, count=20)
            hb = next((m for m in msgs if m[1].get("type") == "heartbeat"), None)
            source = "telemetry"
            if not hb:
                msgs = client.xrevrange(results_stream, count=20)
                hb = next((m for m in msgs if m[1].get("type") == "heartbeat"), None)
                source = "results"
            if not hb:
                return {"status": "unknown", "message": "No heartbeat found in telemetry or results"}
            return {
                "status": hb[1].get("status"),
                "instance_id": hb[1].get("instance_id"),
                "timestamp": hb[1].get("timestamp"),
                "source_stream": source,
            }
        while True:
            res = fetch()
            if out_format == "json":
                click.echo(json.dumps(res, indent=2))
            else:
                click.echo(
                    f"Status: {res.get('status')}  Instance: {res.get('instance_id')}  "
                    f"Time: {res.get('timestamp')}  Stream: {res.get('source_stream')}"
                )
            if not watch:
                break
            time.sleep(5)
    except Exception as exc:
        click.echo(f"Health check failed: {exc}", err=True)
        sys.exit(1)


@cli.command()
def version():
    """Show version information."""
    click.echo(f"Agent CLI version {__version__}")


@cli.command()
@click.option("--log-dir", type=click.Path(), help="Directory where logs are written")
@click.option("--lines", default=100, help="Number of lines to show")
def logs(log_dir, lines):
    """Show recent logs from log directory."""
    path = get_log_file_path(log_dir)
    if not path.exists():
        click.echo(f"No log file found at {path}. Ensure agent started with --log-dir.", err=True)
        sys.exit(1)
    data = path.read_text().splitlines()
    for line in data[-lines:]:
        click.echo(line)


@cli.group()
def debug():
    """Debug utilities."""
    pass


@debug.command("redis")
def debug_redis():
    """Test Redis connection."""
    try:
        ctx = CLIContext()
        client = ctx.build_redis_client()
        pong = client.ping()
        click.echo(f"Redis OK (ping: {pong})")
    except Exception as exc:
        click.echo(f"Redis test failed: {exc}", err=True)
        sys.exit(1)


@debug.command("decode-token")
@click.argument("token", required=False)
@click.option("--show-secrets", is_flag=True, help="Show unredacted secret values")
def debug_decode_token(token, show_secrets):
    """Decode and inspect connection token."""
    try:
        token = token or os.getenv("AGENT_CONNECTION_TOKEN")
        decoded = decode_connection_token(token)
        if not show_secrets:
            decoded = {k: (redact(str(v)) if is_sensitive_key(k) else v) for k, v in decoded.items()}
        click.echo(json.dumps(decoded, indent=2))
    except Exception as exc:
        click.echo(f"Decode failed: {exc}", err=True)
        sys.exit(1)


@debug.command("streams")
@click.option("--show-messages", type=int, default=5, help="Number of recent messages to show")
def debug_streams(show_messages):
    """Inspect Redis streams."""
    try:
        ctx = CLIContext()
        ctx.require_redis_and_key()
        client = ctx.build_redis_client()
        tasks_stream = f"agent:{ctx.agent_key}:tasks"
        results_stream = f"agent:{ctx.agent_key}:results"
        telemetry_stream = f"agent:{ctx.agent_key}:telemetry"
        info_tasks = client.xinfo_stream(tasks_stream)
        info_results = client.xinfo_stream(results_stream)
        info_telemetry = client.xinfo_stream(telemetry_stream)
        click.echo("Streams info:")
        click.echo(
            json.dumps(
                {"tasks": info_tasks, "results": info_results, "telemetry": info_telemetry},
                indent=2,
                default=str,
            )
        )

        recent_tasks = client.xrevrange(tasks_stream, count=show_messages)
        recent_results = client.xrevrange(results_stream, count=show_messages)
        recent_telemetry = client.xrevrange(telemetry_stream, count=show_messages)

        click.echo(f"\nRecent tasks messages ({show_messages}):")
        for mid, data in recent_tasks:
            click.echo(f"{mid}: {data}")

        click.echo(f"\nRecent results messages ({show_messages}):")
        for mid, data in recent_results:
            click.echo(f"{mid}: {data}")

        click.echo(f"\nRecent telemetry messages ({show_messages}):")
        for mid, data in recent_telemetry:
            click.echo(f"{mid}: {data}")
    except Exception as exc:
        click.echo(f"Streams debug failed: {exc}", err=True)
        sys.exit(1)


if __name__ == "__main__":  # pragma: no cover
    cli()
