from __future__ import annotations

import json
import threading
import time
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Optional

from .metrics import MetricsCollector
from .models import AgentMetadata, ComponentHealth, ComponentStatus, TaskMetrics


class HeartbeatSender:
    """Background heartbeat loop."""

    def __init__(self, service: "ObservabilityService", redis_client: Any = None, interval: int = 30):
        self.service = service
        self.interval = interval
        self._thread: Optional[threading.Thread] = None
        self._stop_event = threading.Event()

    def start(self) -> None:
        if self._thread and self._thread.is_alive():
            return
        self._stop_event.clear()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._stop_event.set()
        if self._thread:
            self._thread.join(timeout=2)

    def _run(self) -> None:
        while not self._stop_event.is_set():
            payload = self.service._build_heartbeat()
            try:
                self.service._publish(payload)
            except Exception:
                # Swallow errors to keep heartbeat loop alive; _publish logs upstream.
                pass
            if self.service.telemetry_degraded:
                # Back off when degraded, but still attempt on next cycle
                self._stop_event.wait(self.interval * 2)
            else:
                self._stop_event.wait(self.interval)


class ObservabilityService:
    def __init__(
        self,
        instance_id: str,
        instance_name: Optional[str],
        version: str,
        redis_client: Any = None,
        results_stream: Optional[str] = None,
        telemetry_stream: Optional[str] = None,
        *,
        redis_connection: Any = None,
    ):
        self.metadata = AgentMetadata.build(instance_id, instance_name, version)
        self.component_health: Dict[str, ComponentHealth] = {}
        self.metrics = MetricsCollector()
        # Store RedisConnection reference so we always use the live client
        # after reconnects. Falls back to a raw client for backward compat.
        self._redis_conn = redis_connection
        self._redis_client_fallback = redis_client if redis_connection is None else None
        self.results_stream = results_stream
        self.telemetry_stream = telemetry_stream or results_stream
        self.heartbeat_sender = HeartbeatSender(self, None)
        # Use re-entrant lock to allow nested health registration/update flows
        self._lock = threading.RLock()
        # Telemetry component tracks publish health
        self.register_component("telemetry")
        self._config_version: Optional[int] = None
        self.telemetry_degraded: bool = False

    @property
    def redis(self):
        """Return the current live Redis client, following reconnects."""
        if self._redis_conn is not None:
            return self._redis_conn.client
        return self._redis_client_fallback

    # Component health tracking
    def register_component(self, name: str) -> None:
        with self._lock:
            if name in self.component_health:
                return
            self.component_health[name] = ComponentHealth(
                component=name,
                status=ComponentStatus.UNKNOWN,
                last_check=datetime.now(timezone.utc),
                last_success=None,
                consecutive_failures=0,
                details={},
                error=None,
            )

    def update_component_health(
        self,
        component: str,
        status: ComponentStatus,
        details: Dict[str, Any],
        error: Optional[str] = None,
    ) -> None:
        now = datetime.now(timezone.utc)
        with self._lock:
            if component not in self.component_health:
                self.register_component(component)
            health = self.component_health[component]
            if status == ComponentStatus.HEALTHY:
                health.consecutive_failures = 0
                health.last_success = now
            else:
                health.consecutive_failures += 1
            health.status = status
            health.last_check = now
            health.details = details
            health.error = error

    def get_component_health(self, component: str) -> ComponentHealth:
        with self._lock:
            if component not in self.component_health:
                raise KeyError(f"Component '{component}' not registered")
            return self.component_health[component]

    def get_overall_status(self) -> ComponentStatus:
        with self._lock:
            statuses = [h.status for h in self.component_health.values()]
        if not statuses:
            return ComponentStatus.UNKNOWN
        # Order of severity
        if ComponentStatus.UNHEALTHY in statuses:
            return ComponentStatus.UNHEALTHY
        if ComponentStatus.DEGRADED in statuses:
            return ComponentStatus.DEGRADED
        if all(s == ComponentStatus.UNKNOWN for s in statuses):
            return ComponentStatus.UNKNOWN
        return ComponentStatus.HEALTHY

    # Metrics tracking
    def record_task_start(self, task_id: str, task_type: str) -> None:
        self.metrics.record_start(task_id, task_type)

    def record_task_success(self, task_id: str, execution_time_ms: int) -> None:
        self.metrics.record_success(task_id, execution_time_ms)

    def record_task_failure(self, task_id: str, error_type: str, error_message: str) -> None:
        self.metrics.record_failure(task_id, error_type, error_message)

    def get_metrics(self, window: timedelta = timedelta(hours=24)) -> TaskMetrics:
        return self.metrics.snapshot(window)

    # Reporting helpers
    def generate_health_report(self) -> Dict[str, Any]:
        with self._lock:
            components = {
                name: {
                    "status": health.status.value,
                    "details": health.details,
                    "error": health.error,
                    "last_check": health.last_check.isoformat() + "Z",
                    "last_success": health.last_success.isoformat() + "Z" if health.last_success else None,
                    "consecutive_failures": health.consecutive_failures,
                }
                for name, health in self.component_health.items()
            }
        report = {
            "type": "validation_report",
            "instance_id": self.metadata.instance_id,
            "instance_name": self.metadata.instance_name,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "version": self.metadata.version,
            "config_version": self._config_version,
            "telemetry_degraded": self.telemetry_degraded,
            "metadata": {
                "hostname": self.metadata.hostname,
                "platform": self.metadata.platform,
                "python_version": self.metadata.python_version,
            },
            "components": components,
        }
        return report

    def _build_heartbeat(self) -> Dict[str, Any]:
        uptime_seconds = int((datetime.now(timezone.utc) - self.metadata.started_at).total_seconds())
        with self._lock:
            components = {name: health.status.value for name, health in self.component_health.items()}
        return {
            "type": "heartbeat",
            "instance_id": self.metadata.instance_id,
            "instance_name": self.metadata.instance_name,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "version": self.metadata.version,
            "config_version": self._config_version,
            "telemetry_degraded": self.telemetry_degraded,
            "uptime_seconds": uptime_seconds,
            "status": self.get_overall_status().value,
            "components": components,
        }

    def _build_metrics_report(self, window: timedelta) -> Dict[str, Any]:
        metrics = self.get_metrics(window)
        return {
            "type": "metrics_report",
            "instance_id": self.metadata.instance_id,
            "instance_name": self.metadata.instance_name,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "window_seconds": int(window.total_seconds()),
            "metrics": {
                "total_tasks": metrics.total_tasks,
                "successful_tasks": metrics.successful_tasks,
                "failed_tasks": metrics.failed_tasks,
                "success_rate": (metrics.successful_tasks / metrics.total_tasks) if metrics.total_tasks else 0.0,
                "errors_by_type": {
                    "timeout": metrics.timeout_errors,
                    "llm_error": metrics.llm_errors,
                    "db_error": metrics.db_errors,
                    "desensitization_error": metrics.desensitization_errors,
                },
                "latency": {
                    "avg_ms": metrics.avg_execution_time_ms,
                    "p95_ms": metrics.p95_execution_time_ms,
                    "p99_ms": metrics.p99_execution_time_ms,
                },
            },
        }

    def send_validation_report(self) -> None:
        report = self.generate_health_report()
        self._publish(report)

    def send_metrics_report(self, window: timedelta = timedelta(hours=24)) -> None:
        if self.telemetry_degraded:
            # Back off metrics when degraded
            self._publish_retry_backoff()
            return
        report = self._build_metrics_report(window)
        self._publish(report)

    def publish_config_applied(self, version: Optional[int]) -> None:
        payload = {
            "type": "config_applied",
            "instance_id": self.metadata.instance_id,
            "instance_name": self.metadata.instance_name,
            "config_version": version if version is not None else "",
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
        self._publish(payload)

    def publish_entitlement_error(
        self,
        entitlement: str,
        error: str,
    ) -> None:
        payload = {
            "type": "entitlement_error",
            "instance_id": self.metadata.instance_id,
            "instance_name": self.metadata.instance_name,
            "entitlement": entitlement,
            "error": error,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
        self._publish(payload)

    def publish_entitlements_ok(self, entitlements: list[str]) -> None:
        payload = {
            "type": "entitlements_ok",
            "instance_id": self.metadata.instance_id,
            "instance_name": self.metadata.instance_name,
            "entitlements": entitlements,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
        self._publish(payload)

    def _publish_retry_backoff(self) -> None:
        import time

        backoff = 2
        for _ in range(3):
            try:
                self._publish({"type": "telemetry_probe", "timestamp": datetime.now(timezone.utc).isoformat()})
                self.telemetry_degraded = False
                self.update_component_health("telemetry", ComponentStatus.HEALTHY, {"config_version": self._config_version})
                return
            except Exception:
                time.sleep(backoff)
                backoff = min(backoff * 2, 30)

    def start_heartbeat(self, interval: int = 30) -> None:
        self.heartbeat_sender.interval = interval
        self.heartbeat_sender.start()

    def stop_heartbeat(self) -> None:
        self.heartbeat_sender.stop()

    # Internal publish abstraction
    def _publish(self, payload: Dict[str, Any]) -> None:
        if not self.redis:
            return
        try:
            if self.telemetry_stream and hasattr(self.redis, "xadd"):
                safe_payload = {}
                for k, v in payload.items():
                    if v is None:
                        safe_payload[k] = ""
                    elif isinstance(v, (dict, list)):
                        safe_payload[k] = json.dumps(v)
                    elif isinstance(v, bool):
                        safe_payload[k] = "true" if v else "false"
                    else:
                        safe_payload[k] = v
                self.redis.xadd(self.telemetry_stream, safe_payload)
            else:
                # Fallback to pubsub if no stream configured
                if hasattr(self.redis, "publish"):
                    self.redis.publish("actor:events", json.dumps(payload))
            # success clears degraded flag
            self.telemetry_degraded = False
            self.update_component_health("telemetry", ComponentStatus.HEALTHY, {"config_version": self._config_version})
        except AttributeError:
            # Fallback to JSON string if client expects str
            if hasattr(self.redis, "publish"):
                self.redis.publish("actor:events", json.dumps(payload))
        except Exception as exc:
            import logging

            logging.getLogger(__name__).error("Observability publish failed: %s", exc)
            # Mark telemetry degraded
            self.update_component_health("telemetry", ComponentStatus.DEGRADED, {"last_error": str(exc)})
            self.telemetry_degraded = True
            if hasattr(self.redis, "publish"):
                try:
                    self.redis.publish("actor:events", json.dumps(payload))
                except Exception:
                    pass

    def set_config_version(self, version: Optional[int]) -> None:
        with self._lock:
            self._config_version = version
            self.update_component_health("telemetry", ComponentStatus.HEALTHY, {"config_version": version})
