"""Hot/cold archive tiering for otelctx telemetry data."""

from __future__ import annotations

import hashlib
import json
import logging
import os
import threading
from datetime import UTC, date, datetime, timedelta
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

try:
    import zstandard as zstd
    _ZSTD_AVAILABLE = True
except ImportError:
    _ZSTD_AVAILABLE = False


class ArchiveError(Exception):
    pass


def _compress(data: bytes, level: int = 3) -> bytes:
    if _ZSTD_AVAILABLE:
        cctx = zstd.ZstdCompressor(level=level)
        return cctx.compress(data)
    import gzip
    return gzip.compress(data, compresslevel=level)


def _decompress(data: bytes) -> bytes:
    if _ZSTD_AVAILABLE:
        dctx = zstd.ZstdDecompressor()
        return dctx.decompress(data)
    import gzip
    return gzip.decompress(data)


def _ext() -> str:
    return ".zst" if _ZSTD_AVAILABLE else ".gz"


def _sha256(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


class DailyArchive:
    """Write-once compressed JSONL archive for a single day."""

    def __init__(self, archive_dir: Path, day: date, signal: str = "traces") -> None:
        self.archive_dir = archive_dir
        self.day = day
        self.signal = signal
        self._lock = threading.RLock()
        self._path = archive_dir / f"{day.isoformat()}-{signal}{_ext()}"
        self._manifest_path = archive_dir / f"{day.isoformat()}-{signal}.manifest.json"
        archive_dir.mkdir(parents=True, exist_ok=True)

    def append(self, records: list[dict[str, Any]]) -> None:
        """Compress and append records as JSONL."""
        if not records:
            return
        lines = "\n".join(json.dumps(r, default=str) for r in records) + "\n"
        raw = lines.encode()
        compressed = _compress(raw)
        digest = _sha256(compressed)
        with self._lock:
            mode = "ab" if self._path.exists() else "wb"
            with open(self._path, mode) as f:
                f.write(compressed)
            self._update_manifest(len(records), len(compressed), digest)

    def _update_manifest(self, count: int, compressed_bytes: int, digest: str) -> None:
        if self._manifest_path.exists():
            with open(self._manifest_path) as f:
                manifest = json.load(f)
        else:
            manifest = {
                "day": self.day.isoformat(),
                "signal": self.signal,
                "file": self._path.name,
                "record_count": 0,
                "compressed_bytes": 0,
                "chunks": [],
            }
        manifest["record_count"] += count
        manifest["compressed_bytes"] += compressed_bytes
        manifest["chunks"].append({"digest": digest, "count": count, "bytes": compressed_bytes})
        with open(self._manifest_path, "w") as f:
            json.dump(manifest, f, indent=2)

    def read_all(self) -> list[dict]:
        """Read all records from this archive."""
        if not self._path.exists():
            return []
        # Archive may be multiple concatenated compressed chunks; try single decompress first
        with open(self._path, "rb") as f:
            raw_compressed = f.read()
        # Try reading as a stream of compressed chunks
        records = []
        offset = 0
        data = raw_compressed
        try:
            if _ZSTD_AVAILABLE:
                dctx = zstd.ZstdDecompressor()
                # zstd streaming decompressor handles concatenated frames
                stream_reader = dctx.stream_reader(data)
                decompressed = stream_reader.read()
            else:
                import gzip
                decompressed = gzip.decompress(data)
        except Exception:
            decompressed = b""
        for line in decompressed.decode(errors="replace").splitlines():
            line = line.strip()
            if line:
                try:
                    records.append(json.loads(line))
                except json.JSONDecodeError:
                    pass
        return records

    def search(self, query: str | None, limit: int = 100) -> list[dict]:
        """Linear scan search over archive records."""
        records = self.read_all()
        if not query:
            return records[:limit]
        q = query.lower()
        return [r for r in records if q in json.dumps(r, default=str).lower()][:limit]

    def manifest(self) -> dict | None:
        if not self._manifest_path.exists():
            return None
        with open(self._manifest_path) as f:
            return json.load(f)


class ArchiveTier:
    """Manages hot→cold tiering of telemetry data."""

    def __init__(self, archive_dir: str | Path, hot_days: int = 7) -> None:
        self.archive_dir = Path(archive_dir)
        self.hot_days = hot_days
        self._lock = threading.RLock()
        self.archive_dir.mkdir(parents=True, exist_ok=True)

    def _archive_for(self, day: date, signal: str) -> DailyArchive:
        return DailyArchive(self.archive_dir, day, signal)

    def archive(self, records: list[dict], signal: str = "traces", day: date | None = None) -> None:
        _day = day or date.today()
        archive = self._archive_for(_day, signal)
        archive.append(records)

    def search_cold(
        self,
        query: str | None,
        signal: str = "traces",
        days_back: int = 30,
        limit: int = 100,
    ) -> list[dict]:
        from datetime import timedelta
        results = []
        today = date.today()
        for i in range(self.hot_days, days_back + 1):
            if len(results) >= limit:
                break
            day = today - timedelta(days=i)
            archive = self._archive_for(day, signal)
            hits = archive.search(query, limit=limit - len(results))
            results.extend(hits)
        return results[:limit]

    def list_archives(self, signal: str | None = None) -> list[dict]:
        manifests = []
        pattern = f"*{signal}*" if signal else "*.manifest.json"
        for p in sorted(self.archive_dir.glob("*.manifest.json")):
            if signal and f"-{signal}.manifest.json" not in p.name:
                continue
            try:
                with open(p) as f:
                    manifests.append(json.load(f))
            except Exception:
                pass
        return manifests

    def storage_stats(self) -> dict:
        total_bytes = 0
        archive_count = 0
        record_count = 0
        for p in self.archive_dir.glob("*.manifest.json"):
            try:
                with open(p) as f:
                    m = json.load(f)
                total_bytes += m.get("compressed_bytes", 0)
                record_count += m.get("record_count", 0)
                archive_count += 1
            except Exception:
                pass
        return {
            "archive_dir": str(self.archive_dir),
            "archive_count": archive_count,
            "record_count": record_count,
            "compressed_bytes": total_bytes,
            "compressed_mb": round(total_bytes / (1024 * 1024), 3),
        }


class ArchiveScheduler:
    """Background thread that periodically moves aged-out hot data to cold archive.

    ``fetch_fn(signal, cutoff_dt) -> list[dict]`` must return records older
    than *cutoff_dt* for the given signal type.
    ``delete_fn(signal, cutoff_dt)`` is called after archiving to prune the
    hot store.
    """

    SIGNALS = ("traces", "logs", "metrics")

    def __init__(
        self,
        tier: ArchiveTier,
        fetch_fn,
        delete_fn,
        *,
        hot_days: int = 7,
        check_interval_seconds: float = 86_400.0,  # 24 h
    ) -> None:
        self._tier = tier
        self._fetch_fn = fetch_fn
        self._delete_fn = delete_fn
        self._hot_days = hot_days
        self._interval = check_interval_seconds
        self._stop_event = threading.Event()
        self._thread: threading.Thread | None = None

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

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

    def stop(self, timeout: float = 5.0) -> None:
        self._stop_event.set()
        if self._thread is not None:
            self._thread.join(timeout=timeout)

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _run(self) -> None:
        # Run immediately on first start, then wait for interval
        self._rotate_once()
        while not self._stop_event.wait(timeout=self._interval):
            self._rotate_once()

    def _rotate_once(self) -> None:
        from datetime import timedelta

        cutoff = datetime.now(UTC) - timedelta(days=self._hot_days)
        cutoff_date = cutoff.date()

        for signal in self.SIGNALS:
            try:
                records = self._fetch_fn(signal, cutoff)
                if not records:
                    continue
                # Group by day
                by_day: dict[date, list[dict]] = {}
                for record in records:
                    day = self._extract_day(record, cutoff_date)
                    by_day.setdefault(day, []).append(record)
                for day, day_records in by_day.items():
                    self._tier.archive(day_records, signal=signal, day=day)
                self._delete_fn(signal, cutoff)
            except Exception:  # noqa: BLE001
                pass

    @staticmethod
    def _extract_day(record: dict, fallback: date) -> date:
        """Best-effort extraction of record date for archive bucketing."""
        for key in ("timestamp", "start_time", "time_unix_nano", "created_at"):
            val = record.get(key)
            if val is None:
                continue
            try:
                if isinstance(val, (int, float)):
                    # Nanoseconds if > 1e15, else seconds
                    ts = val / 1e9 if val > 1e15 else float(val)
                    return datetime.fromtimestamp(ts, tz=UTC).date()
                if isinstance(val, str):
                    return datetime.fromisoformat(val.replace("Z", "+00:00")).date()
            except Exception:
                continue
        return fallback


# ----------------------------------------------------------------------
# Archive hardening helpers
# ----------------------------------------------------------------------

def auto_archive(repo, settings) -> dict:
    """Move data older than *hot_days* to cold archive; delete data older than *cold_days*.

    ``repo`` must expose ``session_factory``, and the session must give access
    to the ORM models (Trace, LogEntry, MetricPoint).

    Returns a summary dict with archived/deleted counts per signal.
    """
    from otelctx.db.models import LogEntry, MetricPoint, Span, Trace

    now = datetime.now(UTC)
    hot_cutoff = now - timedelta(days=settings.hot_days)
    cold_cutoff = now - timedelta(days=settings.cold_days)

    tier = ArchiveTier(archive_dir=settings.archive_dir, hot_days=settings.hot_days)
    summary: dict[str, dict[str, int]] = {}

    signal_map: list[tuple[str, type, str]] = [
        ("traces", Trace, "created_at"),
        ("logs", LogEntry, "created_at"),
        ("metrics", MetricPoint, "created_at"),
    ]

    session = repo.session_factory()
    try:
        for signal_name, model, ts_col in signal_map:
            ts_attr = getattr(model, ts_col)

            # --- archive: hot_cutoff < age < cold_cutoff → move to cold ---
            to_archive = (
                session.query(model)
                .filter(ts_attr < hot_cutoff, ts_attr >= cold_cutoff)
                .all()
            )
            archived_count = 0
            if to_archive:
                records = []
                for row in to_archive:
                    try:
                        d = {c.name: getattr(row, c.name) for c in model.__table__.columns}
                        records.append(d)
                    except Exception:
                        pass
                if records:
                    tier.archive(records, signal=signal_name)
                    archived_count = len(records)

            # --- delete: older than cold_cutoff ---
            if signal_name == "traces":
                # Delete child spans first to respect FK constraint
                trace_ids = [
                    r.id
                    for r in session.query(Trace.id).filter(ts_attr < cold_cutoff).all()
                ]
                deleted_count = 0
                if trace_ids:
                    session.query(Span).filter(Span.trace_id.in_(trace_ids)).delete(
                        synchronize_session=False
                    )
                    deleted_count = (
                        session.query(model).filter(ts_attr < cold_cutoff).delete(
                            synchronize_session=False
                        )
                    )
            else:
                deleted_count = (
                    session.query(model)
                    .filter(ts_attr < cold_cutoff)
                    .delete(synchronize_session=False)
                )

            summary[signal_name] = {
                "archived": archived_count,
                "deleted": deleted_count,
            }

        session.commit()
    except Exception:
        session.rollback()
        logger.exception("auto_archive failed")
        raise
    finally:
        session.close()

    return summary


def compact_cold_storage(archive_dir: str | Path) -> dict:
    """Aggregate old metrics in archive: keep hourly averages, drop debug logs.

    Rewrites archive files in-place.  Returns a summary of actions taken.
    """
    archive_dir = Path(archive_dir)
    summary = {"compacted_files": 0, "dropped_debug_logs": 0}

    for manifest_path in sorted(archive_dir.glob("*.manifest.json")):
        try:
            with open(manifest_path) as f:
                manifest = json.load(f)
        except Exception:
            continue

        signal = manifest.get("signal", "")
        archive_file = archive_dir / manifest.get("file", "")
        if not archive_file.exists():
            continue

        archive = DailyArchive(archive_dir, date.fromisoformat(manifest["day"]), signal)
        records = archive.read_all()
        if not records:
            continue

        rewrite = False

        if signal == "metrics":
            # Group by metric_name + hour and keep hourly averages
            from collections import defaultdict

            hourly: dict[str, dict[int, list[float]]] = defaultdict(lambda: defaultdict(list))
            for r in records:
                name = r.get("metric_name", "unknown")
                ts_val = r.get("time") or r.get("timestamp") or r.get("created_at")
                try:
                    if isinstance(ts_val, str):
                        dt = datetime.fromisoformat(ts_val.replace("Z", "+00:00"))
                    elif isinstance(ts_val, (int, float)):
                        dt = datetime.fromtimestamp(ts_val if ts_val < 1e15 else ts_val / 1e9, tz=UTC)
                    else:
                        continue
                except Exception:
                    continue
                hour_key = int(dt.timestamp()) // 3600
                hourly[name][hour_key].append(float(r.get("value", 0)))

            compacted: list[dict] = []
            for name, hours in hourly.items():
                for hour_key, values in sorted(hours.items()):
                    compacted.append({
                        "metric_name": name,
                        "time": datetime.fromtimestamp(hour_key * 3600, tz=UTC).isoformat(),
                        "value": sum(values) / len(values),
                        "count": len(values),
                        "min": min(values),
                        "max": max(values),
                        "aggregation": "hourly_avg",
                    })
            records = compacted
            rewrite = True

        elif signal == "logs":
            # Drop debug-level logs
            original_len = len(records)
            records = [
                r for r in records
                if r.get("severity_text", "").upper() not in ("DEBUG", "TRACE")
            ]
            dropped = original_len - len(records)
            if dropped > 0:
                summary["dropped_debug_logs"] += dropped
                rewrite = True

        if rewrite and records:
            # Overwrite the archive file
            lines = "\n".join(json.dumps(r, default=str) for r in records) + "\n"
            raw = lines.encode()
            compressed = _compress(raw)
            digest = _sha256(compressed)
            with open(archive_file, "wb") as f:
                f.write(compressed)
            # Update manifest
            manifest["record_count"] = len(records)
            manifest["compressed_bytes"] = len(compressed)
            manifest["chunks"] = [{"digest": digest, "count": len(records), "bytes": len(compressed)}]
            manifest["compacted"] = True
            with open(manifest_path, "w") as f:
                json.dump(manifest, f, indent=2)
            summary["compacted_files"] += 1

    return summary


def get_retention_stats(repo, settings) -> dict:
    """Return retention statistics: hot_count, cold_count, oldest_hot, oldest_cold."""
    from sqlalchemy import func as sa_func

    from otelctx.db.models import LogEntry, MetricPoint, Trace

    now = datetime.now(UTC)
    hot_cutoff = now - timedelta(days=settings.hot_days)

    session = repo.session_factory()
    try:
        hot_count = 0
        cold_count = 0
        oldest_hot: datetime | None = None
        oldest_cold: datetime | None = None

        for model, ts_col in [(Trace, "created_at"), (LogEntry, "created_at"), (MetricPoint, "created_at")]:
            ts_attr = getattr(model, ts_col)

            hc = session.query(sa_func.count(model.id)).filter(ts_attr >= hot_cutoff).scalar() or 0
            hot_count += hc

            cc = session.query(sa_func.count(model.id)).filter(ts_attr < hot_cutoff).scalar() or 0
            cold_count += cc

            oh = session.query(sa_func.min(ts_attr)).filter(ts_attr >= hot_cutoff).scalar()
            if oh and (oldest_hot is None or oh < oldest_hot):
                oldest_hot = oh

            oc = session.query(sa_func.min(ts_attr)).filter(ts_attr < hot_cutoff).scalar()
            if oc and (oldest_cold is None or oc < oldest_cold):
                oldest_cold = oc

        return {
            "hot_count": hot_count,
            "cold_count": cold_count,
            "oldest_hot": oldest_hot.isoformat() if oldest_hot else None,
            "oldest_cold": oldest_cold.isoformat() if oldest_cold else None,
            "hot_days": settings.hot_days,
            "cold_days": settings.cold_days,
        }
    finally:
        session.close()
