"""
Brynq Runtime — Execution Metrics Reporter.

Collects per-step and per-plan execution metadata and sends it to the Brynq
Cloud for quality improvement. The reporter NEVER transmits user content,
model outputs, prompt text, or API keys — all outbound data passes through
the privacy filter (``runtime.reporter.privacy``) before leaving the machine.

Batching strategy:
    - Reports are queued in memory as they arrive.
    - A flush happens on whichever comes FIRST:
        a) Chain completion (``report_execution()`` triggers immediate flush)
        b) Timer expiry (every ``BATCH_INTERVAL_SECONDS``, default 30s)
        c) Queue age limit (``MAX_QUEUE_AGE_SECONDS``, default 60s) — if
           reports have been sitting in the queue for 60+ seconds without a
           flush, one is triggered automatically regardless of the timer
    - If the cloud is unreachable, reports are persisted to a local JSONL file
      (``~/.brynq/pending_reports.jsonl``) and retried on the next flush cycle.

Failure handling:
    - Network errors are caught and logged, never raised to the caller.
    - The local retry queue is capped at ``MAX_PENDING_REPORTS`` to prevent
      unbounded disk usage.
    - On successful reconnection, pending reports are drained oldest-first.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

from .privacy import SanitizedReport, sanitize_report, to_wire_dict

log = logging.getLogger("brynq.runtime.metrics")

# ── Configuration ──────────────────────────────────────────────────────

BATCH_INTERVAL_SECONDS: float = 30.0
MAX_QUEUE_AGE_SECONDS: float = 60.0
MAX_BATCH_SIZE: int = 100
MAX_PENDING_REPORTS: int = 1000
DEFAULT_CLOUD_URL: str = "https://api.brynq.ai/v1/report"
PENDING_REPORTS_DIR: str = os.path.join(
    os.path.expanduser("~"), ".brynq"
)
PENDING_REPORTS_FILE: str = os.path.join(
    PENDING_REPORTS_DIR, "pending_reports.jsonl"
)
RETRY_QUEUE_FILE: str = os.path.join(
    PENDING_REPORTS_DIR, "metrics_retry.jsonl"
)
MAX_RETRY_QUEUE_ENTRIES: int = 500


# ── Data models ────────────────────────────────────────────────────────

@dataclass
class StepMetrics:
    """Raw metrics collected for a single execution step.

    These are gathered by the plan executor during step execution.
    They may contain identifiers that reference user content — the
    privacy filter strips those before transmission.

    Attributes:
        step_id: Identifier matching the PlanStep.
        elapsed_ms: Wall clock time for this step.
        tokens_used: Total tokens consumed (input + output).
        success: Whether the step completed without error.
        error_type: Categorical error label (e.g., "timeout", "rate_limit",
                    "model_error", "auth_error"). None if success.
        model_name: Which model executed this step.
        backend: Execution backend ("ollama", "claude", "gemini", etc.).
        task_type: Task classification label.
    """

    step_id: str
    elapsed_ms: float
    tokens_used: int
    success: bool
    error_type: str | None = None
    model_name: str = ""
    backend: str = ""
    task_type: str = ""


@dataclass
class PlanMetrics:
    """Aggregated metrics for a complete plan execution.

    Computed from the individual StepMetrics after the plan finishes.

    Attributes:
        plan_id: Which execution plan was run.
        total_elapsed_ms: End-to-end wall clock time.
        steps_completed: Number of steps that succeeded.
        steps_failed: Number of steps that failed.
        step_metrics: Per-step raw metrics.
        reported_at: Unix timestamp when metrics were collected.
    """

    plan_id: str
    total_elapsed_ms: float
    steps_completed: int
    steps_failed: int
    step_metrics: list[StepMetrics] = field(default_factory=list)
    reported_at: float = field(default_factory=time.time)


# ── Transport interface ────────────────────────────────────────────────

class ReportTransport:
    """Abstract transport for sending reports to the cloud.

    The default implementation uses HTTP POST. Tests can substitute
    a mock transport.
    """

    async def send(self, reports: list[dict[str, Any]]) -> bool:
        """Send a batch of sanitized report dicts to the cloud.

        Args:
            reports: List of sanitized report dicts (output of ``to_wire_dict``).

        Returns:
            True if the cloud accepted the batch, False on any failure.
        """
        raise NotImplementedError


class HttpTransport(ReportTransport):
    """HTTP POST transport targeting the Brynq Cloud report endpoint.

    Uses ``aiohttp`` if available, falls back to ``urllib`` in a thread.
    """

    def __init__(self, url: str = DEFAULT_CLOUD_URL) -> None:
        self.url = url

    async def send(self, reports: list[dict[str, Any]]) -> bool:
        """POST a JSON batch of reports to the cloud endpoint.

        Args:
            reports: Sanitized report dicts.

        Returns:
            True on 2xx response, False otherwise.
        """
        payload = json.dumps({"reports": reports}).encode("utf-8")

        try:
            # Prefer aiohttp if available
            import aiohttp

            async with aiohttp.ClientSession() as session:
                async with session.post(
                    self.url,
                    data=payload,
                    headers={"Content-Type": "application/json"},
                    timeout=aiohttp.ClientTimeout(total=10),
                ) as resp:
                    return 200 <= resp.status < 300
        except ImportError:
            # Fallback: urllib in a thread pool
            return await self._send_urllib(payload)
        except Exception:
            log.warning("Failed to send reports via aiohttp", exc_info=True)
            return False

    async def _send_urllib(self, payload: bytes) -> bool:
        """Blocking urllib fallback, run in the default executor."""
        import urllib.request

        def _do_post() -> bool:
            req = urllib.request.Request(
                self.url,
                data=payload,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            try:
                with urllib.request.urlopen(req, timeout=10) as resp:
                    return 200 <= resp.status < 300
            except Exception:
                log.warning("Failed to send reports via urllib", exc_info=True)
                return False

        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, _do_post)


# ── Persistence for retry queue ────────────────────────────────────────

def _ensure_pending_dir() -> None:
    """Create the ~/.brynq directory if it doesn't exist."""
    os.makedirs(PENDING_REPORTS_DIR, exist_ok=True)


def _append_pending(reports: list[dict[str, Any]]) -> int:
    """Append reports to the local retry queue.

    Args:
        reports: Sanitized report dicts to persist.

    Returns:
        Number of reports successfully written.
    """
    _ensure_pending_dir()
    written = 0
    try:
        with open(PENDING_REPORTS_FILE, "a", encoding="utf-8") as f:
            for report in reports:
                f.write(json.dumps(report) + "\n")
                written += 1
    except OSError:
        log.warning("Failed to write pending reports to disk", exc_info=True)
    return written


def _load_pending() -> list[dict[str, Any]]:
    """Load all pending reports from the local retry queue.

    Returns:
        List of report dicts, oldest first. Empty list if file doesn't exist
        or is unreadable.
    """
    if not os.path.exists(PENDING_REPORTS_FILE):
        return []

    reports: list[dict[str, Any]] = []
    try:
        with open(PENDING_REPORTS_FILE, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    try:
                        reports.append(json.loads(line))
                    except json.JSONDecodeError:
                        log.warning("Skipping malformed pending report line")
    except OSError:
        log.warning("Failed to read pending reports from disk", exc_info=True)

    return reports


def _clear_pending() -> None:
    """Remove the pending reports file after successful drain."""
    try:
        if os.path.exists(PENDING_REPORTS_FILE):
            os.remove(PENDING_REPORTS_FILE)
    except OSError:
        log.warning("Failed to clear pending reports file", exc_info=True)


def _truncate_pending(keep: list[dict[str, Any]]) -> None:
    """Rewrite the pending file with only the reports that still need sending.

    Args:
        keep: Reports that failed to send and should be retained.
    """
    _ensure_pending_dir()
    try:
        with open(PENDING_REPORTS_FILE, "w", encoding="utf-8") as f:
            for report in keep[-MAX_PENDING_REPORTS:]:
                f.write(json.dumps(report) + "\n")
    except OSError:
        log.warning("Failed to truncate pending reports file", exc_info=True)


# ── Retry queue persistence (metrics_retry.jsonl) ─────────────────────


def _append_retry_queue(reports: list[dict[str, Any]]) -> int:
    """Append failed reports to the retry queue file for later recovery.

    Each entry is a JSON object with the original report data and failure
    metadata (``failed_at`` timestamp).

    Args:
        reports: Sanitized report dicts that failed to send.

    Returns:
        Number of entries successfully written.
    """
    _ensure_pending_dir()
    written = 0
    failed_at = time.time()
    try:
        with open(RETRY_QUEUE_FILE, "a", encoding="utf-8") as f:
            for report in reports:
                entry = {"failed_at": failed_at, "report": report}
                f.write(json.dumps(entry) + "\n")
                written += 1
        log.debug("Wrote %d failed reports to retry queue", written)
    except OSError:
        log.warning("Failed to write to metrics retry queue", exc_info=True)
    # Enforce cap — trim oldest entries if the file has grown too large
    if written > 0:
        _cap_retry_queue()
    return written


def _load_retry_queue() -> list[dict[str, Any]]:
    """Load all entries from the retry queue file.

    Returns:
        List of report dicts (unwrapped from the retry envelope),
        oldest first. Empty list if file doesn't exist or is unreadable.
    """
    if not os.path.exists(RETRY_QUEUE_FILE):
        return []

    reports: list[dict[str, Any]] = []
    try:
        with open(RETRY_QUEUE_FILE, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    entry = json.loads(line)
                    # Unwrap the retry envelope → just the report
                    if isinstance(entry, dict) and "report" in entry:
                        reports.append(entry["report"])
                    else:
                        # Bare report (forward-compatible)
                        reports.append(entry)
                except json.JSONDecodeError:
                    log.warning("Skipping malformed retry queue line")
    except OSError:
        log.warning("Failed to read metrics retry queue", exc_info=True)

    return reports


def _clear_retry_queue() -> None:
    """Remove the retry queue file after successful recovery."""
    try:
        if os.path.exists(RETRY_QUEUE_FILE):
            os.remove(RETRY_QUEUE_FILE)
    except OSError:
        log.warning("Failed to clear metrics retry queue file", exc_info=True)


def _cap_retry_queue() -> None:
    """Trim the retry queue file to at most MAX_RETRY_QUEUE_ENTRIES."""
    if not os.path.exists(RETRY_QUEUE_FILE):
        return
    try:
        with open(RETRY_QUEUE_FILE, "r", encoding="utf-8") as f:
            lines = f.readlines()
        if len(lines) <= MAX_RETRY_QUEUE_ENTRIES:
            return
        dropped = len(lines) - MAX_RETRY_QUEUE_ENTRIES
        log.warning("Retry queue cap: dropping %d oldest entries", dropped)
        with open(RETRY_QUEUE_FILE, "w", encoding="utf-8") as f:
            f.writelines(lines[-MAX_RETRY_QUEUE_ENTRIES:])
    except OSError:
        log.warning("Failed to cap metrics retry queue", exc_info=True)


# ── Main reporter ──────────────────────────────────────────────────────

class ExecutionReporter:
    """Batched, retry-safe execution metrics reporter.

    Usage::

        reporter = ExecutionReporter()
        await reporter.start()  # Starts the periodic flush timer

        # After a chain completes:
        success = await reporter.report_execution(plan_id, step_results)

        # On shutdown:
        await reporter.stop()

    Attributes:
        transport: The transport used to send reports to the cloud.
        batch_interval: Seconds between periodic flushes.
    """

    def __init__(
        self,
        transport: ReportTransport | None = None,
        batch_interval: float = BATCH_INTERVAL_SECONDS,
    ) -> None:
        self.transport = transport or HttpTransport()
        self.batch_interval = batch_interval
        self._queue: list[SanitizedReport] = []
        self._flush_task: asyncio.Task | None = None
        self._running: bool = False
        self._oldest_queued_at: float | None = None
        self._seen_keys: set[tuple[str, str, float]] = set()

    async def start(self) -> None:
        """Start the periodic flush timer.

        Safe to call multiple times — subsequent calls are no-ops.
        """
        if self._running:
            return
        self._running = True
        self._flush_task = asyncio.create_task(self._periodic_flush())
        log.info(
            "Execution reporter started (flush every %.0fs)", self.batch_interval
        )

    async def stop(self) -> None:
        """Stop the periodic flush and drain remaining reports.

        Performs one final flush attempt. Any unsent reports are persisted
        to the local retry queue.
        """
        self._running = False
        if self._flush_task is not None:
            self._flush_task.cancel()
            try:
                await self._flush_task
            except asyncio.CancelledError:
                pass
            self._flush_task = None

        # Final flush
        await self._flush()
        log.info("Execution reporter stopped")

    async def flush_and_close(self) -> bool:
        """Send all queued reports and shut down the reporter.

        Convenience method for graceful shutdown. Flushes every queued
        report (retrying once on failure) and then stops the periodic
        flush timer. Any reports that still cannot be sent are persisted
        to the local retry queue for the next session.

        Returns:
            True if all reports were sent successfully, False if some
            had to be persisted locally for retry.
        """
        log.info(
            "Graceful shutdown: flushing %d queued reports", len(self._queue)
        )

        # Stop accepting new periodic flushes
        self._running = False
        if self._flush_task is not None:
            self._flush_task.cancel()
            try:
                await self._flush_task
            except asyncio.CancelledError:
                pass
            self._flush_task = None

        # First flush attempt
        success = await self._flush()

        # If first attempt failed and there are pending reports on disk,
        # retry once to give transient network issues a chance to clear
        if not success:
            log.info("First flush failed, retrying once before exit")
            success = await self._flush()

        log.info(
            "Graceful shutdown complete (all_sent=%s)", success
        )
        return success

    async def report_execution(
        self,
        plan_id: str,
        step_results: list[StepMetrics],
    ) -> bool:
        """Report execution results for a completed plan.

        Builds a PlanMetrics from the step results, sanitizes it through
        the privacy filter, queues it, and triggers an immediate flush.

        Args:
            plan_id: The execution plan's unique identifier.
            step_results: Per-step metrics collected during execution.

        Returns:
            True if the report was successfully sent to the cloud (or queued
            locally for retry). False only on catastrophic failure (should
            not happen in practice).
        """
        # Build aggregate metrics
        reported_at = time.time()

        # Deduplicate: skip steps already reported with same plan_id + step_id + timestamp
        new_steps: list[StepMetrics] = []
        for s in step_results:
            key = (plan_id, s.step_id, reported_at)
            if key in self._seen_keys:
                log.debug(
                    "Skipping duplicate report: plan=%s step=%s ts=%f",
                    plan_id, s.step_id, reported_at,
                )
                continue
            self._seen_keys.add(key)
            new_steps.append(s)

        if not new_steps:
            log.debug("All steps in plan %s are duplicates, skipping", plan_id)
            return True

        total_ms = sum(s.elapsed_ms for s in new_steps)
        completed = sum(1 for s in new_steps if s.success)
        failed = sum(1 for s in new_steps if not s.success)

        plan_metrics = PlanMetrics(
            plan_id=plan_id,
            total_elapsed_ms=total_ms,
            steps_completed=completed,
            steps_failed=failed,
            step_metrics=new_steps,
            reported_at=reported_at,
        )

        # Convert to raw dict for sanitization
        raw = _plan_metrics_to_raw(plan_metrics)

        # Sanitize through privacy filter — this is the chokepoint
        sanitized = sanitize_report(raw)

        # Queue for sending — track when first report entered the queue
        if not self._queue:
            self._oldest_queued_at = time.time()
        self._queue.append(sanitized)

        # Immediate flush on chain completion
        sent = await self._flush()
        return sent

    async def _flush(self) -> bool:
        """Send all queued reports to the cloud.

        Also attempts to drain any pending reports from the local retry file.

        Returns:
            True if all reports were sent successfully.
        """
        has_pending = os.path.exists(PENDING_REPORTS_FILE)
        has_retry = os.path.exists(RETRY_QUEUE_FILE)
        if not self._queue and not has_pending and not has_retry:
            return True

        # Convert queued reports to wire dicts
        wire_reports = [to_wire_dict(r) for r in self._queue]
        self._queue.clear()
        self._oldest_queued_at = None

        # Load any previously-failed reports
        pending = _load_pending()
        if pending:
            log.info("Found %d pending reports from previous session", len(pending))

        # Recover from retry queue (metrics_retry.jsonl)
        retried = _load_retry_queue()
        if retried:
            log.info("Recovering %d reports from retry queue", len(retried))

        # Combine: pending first (oldest), then retried, then new
        all_reports = pending + retried + wire_reports

        if not all_reports:
            return True

        # Cap to prevent unbounded growth
        if len(all_reports) > MAX_PENDING_REPORTS:
            dropped = len(all_reports) - MAX_PENDING_REPORTS
            all_reports = all_reports[-MAX_PENDING_REPORTS:]
            log.warning("Dropped %d oldest pending reports (queue cap)", dropped)

        # Enforce max batch size — only send first MAX_BATCH_SIZE, defer the rest
        deferred: list[dict[str, Any]] = []
        if len(all_reports) > MAX_BATCH_SIZE:
            overflow = len(all_reports) - MAX_BATCH_SIZE
            log.warning(
                "Batch overflow: %d reports exceed max batch size %d, "
                "deferring %d to next flush cycle",
                len(all_reports),
                MAX_BATCH_SIZE,
                overflow,
            )
            deferred = all_reports[MAX_BATCH_SIZE:]
            all_reports = all_reports[:MAX_BATCH_SIZE]

        # Attempt to send
        success = await self.transport.send(all_reports)

        if success:
            _clear_pending()
            _clear_retry_queue()
            log.debug("Flushed %d reports to cloud", len(all_reports))
            # Re-persist deferred reports for next cycle
            if deferred:
                _append_pending(deferred)
            return True
        else:
            # Persist for retry (both unsent and deferred)
            _clear_pending()
            _append_pending(all_reports + deferred)
            # Write failed reports to metrics_retry.jsonl for recovery
            _clear_retry_queue()
            _append_retry_queue(all_reports + deferred)
            log.warning(
                "Cloud unreachable — queued %d reports locally for retry",
                len(all_reports) + len(deferred),
            )
            return False

    async def _periodic_flush(self) -> None:
        """Background task: flush every ``batch_interval`` seconds.

        Also triggers an early flush if queued reports have been waiting
        longer than ``MAX_QUEUE_AGE_SECONDS`` (default 60s).
        """
        while self._running:
            try:
                # Sleep in short increments so we can check queue age
                sleep_time = self.batch_interval
                if self._oldest_queued_at is not None:
                    age = time.time() - self._oldest_queued_at
                    remaining = MAX_QUEUE_AGE_SECONDS - age
                    if remaining <= 0:
                        # Queue has aged past the limit — flush now
                        if self._running:
                            log.debug(
                                "Queue age %.1fs exceeds %ds limit, flushing",
                                age, int(MAX_QUEUE_AGE_SECONDS),
                            )
                            await self._flush()
                        continue
                    sleep_time = min(sleep_time, remaining)

                await asyncio.sleep(sleep_time)
                if self._running:
                    # Check age again after sleep in case reports were queued
                    if self._oldest_queued_at is not None:
                        age = time.time() - self._oldest_queued_at
                        if age >= MAX_QUEUE_AGE_SECONDS:
                            log.debug(
                                "Queue age %.1fs exceeds %ds limit, flushing",
                                age, int(MAX_QUEUE_AGE_SECONDS),
                            )
                    await self._flush()
            except asyncio.CancelledError:
                break
            except Exception:
                log.warning("Periodic flush failed", exc_info=True)


# ── Helpers ────────────────────────────────────────────────────────────

def _plan_metrics_to_raw(metrics: PlanMetrics) -> dict[str, Any]:
    """Convert PlanMetrics to the raw dict expected by ``sanitize_report()``.

    This is the bridge between the runtime's internal data model and the
    privacy filter's input format.

    Args:
        metrics: Aggregated plan metrics.

    Returns:
        Dict with plan-level and step-level fields.
    """
    steps = []
    for s in metrics.step_metrics:
        steps.append({
            "step_id": s.step_id,
            "elapsed_ms": s.elapsed_ms,
            "tokens_used": s.tokens_used,
            "success": s.success,
            "error_type": s.error_type,
            "model_name": s.model_name,
            "backend": s.backend,
            "task_type": s.task_type,
        })

    return {
        "plan_id": metrics.plan_id,
        "total_elapsed_ms": metrics.total_elapsed_ms,
        "steps_completed": metrics.steps_completed,
        "steps_failed": metrics.steps_failed,
        "steps": steps,
        "reported_at": metrics.reported_at,
    }
