"""Repository methods for otelctx."""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from datetime import UTC, datetime
from pathlib import Path

from sqlalchemy import case, func, select
from sqlalchemy.orm import selectinload

from common.db import Base, create_engine_with_wal, create_session_factory, session_scope
from otelctx.db.models import LogEntry, MetricPoint, Span, Trace


def _parse_timestamp(value: str | None) -> datetime | None:
    if not value:
        return None
    normalized = value.replace("Z", "+00:00")
    parsed = datetime.fromisoformat(normalized)
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=UTC)
    return parsed.astimezone(UTC)


@dataclass(slots=True)
class SpanRecord:
    span_id: str
    name: str
    service_name: str = "unknown"
    parent_span_id: str = ""
    kind: str = "INTERNAL"
    status_code: str = "UNSET"
    start_time: str = ""
    end_time: str = ""
    attributes: dict[str, object] = field(default_factory=dict)


@dataclass(slots=True)
class TraceRecord:
    trace_id: str
    root_service_name: str = "unknown"
    root_span_name: str = ""
    status_code: str = "UNSET"
    spans: list[SpanRecord] = field(default_factory=list)


@dataclass(slots=True)
class LogRecord:
    timestamp: str = ""
    observed_time: str = ""
    service_name: str = "unknown"
    severity_text: str = "INFO"
    body: str = ""
    trace_id: str = ""
    span_id: str = ""
    attributes: dict[str, object] = field(default_factory=dict)


@dataclass(slots=True)
class MetricRecord:
    service_name: str = "unknown"
    name: str = ""
    metric_type: str = "gauge"
    unit: str = ""
    temporality: str = ""
    start_time: str = ""
    time: str = ""
    value: float = 0.0
    attributes: dict[str, object] = field(default_factory=dict)


@dataclass(slots=True)
class ServiceSummary:
    service_name: str
    span_count: int
    trace_count: int
    error_count: int

    def to_dict(self) -> dict[str, object]:
        return asdict(self)


@dataclass(slots=True)
class OtelStats:
    traces: int = 0
    spans: int = 0
    logs: int = 0
    metrics: int = 0
    services: int = 0
    error_spans: int = 0

    def to_dict(self) -> dict[str, int]:
        return asdict(self)


class OtelRepository:
    """Persistence wrapper for otelctx."""

    def __init__(self, db_path: str | Path, echo: bool = False) -> None:
        self.db_path = Path(db_path)
        self.engine = create_engine_with_wal(f"sqlite:///{self.db_path}", echo=echo)
        self.session_factory = create_session_factory(self.engine)

    def create_schema(self) -> None:
        Base.metadata.create_all(
            self.engine,
            tables=[Trace.__table__, Span.__table__, LogEntry.__table__, MetricPoint.__table__],
        )

    def ingest_traces(self, traces: list[TraceRecord]) -> dict[str, int]:
        with session_scope(self.session_factory) as session:
            ingested_traces = 0
            ingested_spans = 0
            for record in traces:
                existing = session.scalar(
                    select(Trace).options(selectinload(Trace.spans)).where(Trace.trace_id == record.trace_id)
                )
                start_time, end_time, duration_ms = self._trace_timing(record.spans)
                if existing is None:
                    trace_row = Trace(
                        trace_id=record.trace_id,
                        root_service_name=record.root_service_name,
                        root_span_name=record.root_span_name,
                        status_code=record.status_code,
                        span_count=len(record.spans),
                        start_time=start_time,
                        end_time=end_time,
                        duration_ms=duration_ms,
                    )
                    session.add(trace_row)
                    session.flush()
                else:
                    trace_row = existing
                    trace_row.root_service_name = record.root_service_name
                    trace_row.root_span_name = record.root_span_name
                    trace_row.status_code = record.status_code
                    trace_row.span_count = len(record.spans)
                    trace_row.start_time = start_time
                    trace_row.end_time = end_time
                    trace_row.duration_ms = duration_ms
                    for span in list(trace_row.spans):
                        session.delete(span)
                    session.flush()

                for span in record.spans:
                    span_start = _parse_timestamp(span.start_time)
                    span_end = _parse_timestamp(span.end_time)
                    span_duration = 0.0
                    if span_start is not None and span_end is not None:
                        span_duration = max(0.0, (span_end - span_start).total_seconds() * 1000.0)
                    session.add(
                        Span(
                            trace_id=trace_row.id,
                            span_id=span.span_id,
                            parent_span_id=span.parent_span_id or None,
                            name=span.name,
                            service_name=span.service_name,
                            kind=span.kind,
                            status_code=span.status_code,
                            start_time=span_start,
                            end_time=span_end,
                            duration_ms=span_duration,
                            attributes_json=json.dumps(span.attributes or {}, sort_keys=True),
                        )
                    )
                    ingested_spans += 1

                ingested_traces += 1

            return {"traces": ingested_traces, "spans": ingested_spans}

    def ingest_logs(self, logs: list[LogRecord]) -> dict[str, int]:
        with session_scope(self.session_factory) as session:
            ingested_logs = 0
            for record in logs:
                session.add(
                    LogEntry(
                        timestamp=_parse_timestamp(record.timestamp),
                        observed_time=_parse_timestamp(record.observed_time),
                        service_name=record.service_name,
                        severity_text=record.severity_text,
                        body=record.body,
                        trace_id=record.trace_id or None,
                        span_id=record.span_id or None,
                        attributes_json=json.dumps(record.attributes or {}, sort_keys=True),
                    )
                )
                ingested_logs += 1
            return {"logs": ingested_logs}

    def ingest_metrics(self, metrics: list[MetricRecord]) -> dict[str, int]:
        with session_scope(self.session_factory) as session:
            ingested_metrics = 0
            for record in metrics:
                session.add(
                    MetricPoint(
                        service_name=record.service_name,
                        metric_name=record.name,
                        metric_type=record.metric_type,
                        unit=record.unit or None,
                        temporality=record.temporality or None,
                        start_time=_parse_timestamp(record.start_time),
                        time=_parse_timestamp(record.time),
                        value=record.value,
                        attributes_json=json.dumps(record.attributes or {}, sort_keys=True),
                    )
                )
                ingested_metrics += 1
            return {"metrics": ingested_metrics}

    def list_traces(self, limit: int = 20, offset: int = 0) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            rows = list(
                session.scalars(
                    select(Trace).order_by(Trace.start_time.desc().nullslast(), Trace.id.desc()).limit(limit).offset(offset)
                )
            )
            return [self._trace_summary(row) for row in rows]

    def get_trace(self, trace_id: str) -> dict[str, object] | None:
        with session_scope(self.session_factory) as session:
            row = session.scalar(
                select(Trace)
                .options(selectinload(Trace.spans))
                .where(Trace.trace_id == trace_id)
            )
            if row is None:
                return None
            payload = self._trace_summary(row)
            payload["spans"] = [
                {
                    "span_id": span.span_id,
                    "parent_span_id": span.parent_span_id,
                    "name": span.name,
                    "service_name": span.service_name,
                    "kind": span.kind,
                    "status_code": span.status_code,
                    "duration_ms": round(span.duration_ms, 3),
                    "attributes": json.loads(span.attributes_json or "{}"),
                }
                for span in sorted(row.spans, key=lambda item: (item.start_time or datetime.min.replace(tzinfo=UTC), item.id))
            ]
            return payload

    def list_services(self) -> list[ServiceSummary]:
        with session_scope(self.session_factory) as session:
            rows = session.execute(
                select(
                    Span.service_name,
                    func.count(Span.id),
                    func.count(func.distinct(Trace.trace_id)),
                    func.sum(case((Span.status_code == "ERROR", 1), else_=0)),
                )
                .join(Trace, Span.trace_id == Trace.id)
                .group_by(Span.service_name)
                .order_by(func.count(Span.id).desc(), Span.service_name.asc())
            )
            return [
                ServiceSummary(
                    service_name=service_name,
                    span_count=int(span_count or 0),
                    trace_count=int(trace_count or 0),
                    error_count=int(error_count or 0),
                )
                for service_name, span_count, trace_count, error_count in rows
            ]

    def list_logs(self, limit: int = 20, offset: int = 0) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            rows = list(
                session.scalars(
                    select(LogEntry)
                    .order_by(LogEntry.timestamp.desc().nullslast(), LogEntry.id.desc())
                    .limit(limit)
                    .offset(offset)
                )
            )
            return [self._log_summary(row) for row in rows]

    def list_metrics(self, limit: int = 20, offset: int = 0) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            rows = list(
                session.scalars(
                    select(MetricPoint)
                    .order_by(MetricPoint.time.desc().nullslast(), MetricPoint.id.desc())
                    .limit(limit)
                    .offset(offset)
                )
            )
            return [self._metric_summary(row) for row in rows]

    def get_stats(self) -> OtelStats:
        with session_scope(self.session_factory) as session:
            trace_count = session.scalar(select(func.count()).select_from(Trace)) or 0
            span_count = session.scalar(select(func.count()).select_from(Span)) or 0
            log_count = session.scalar(select(func.count()).select_from(LogEntry)) or 0
            metric_count = session.scalar(select(func.count()).select_from(MetricPoint)) or 0
            service_count = session.scalar(select(func.count(func.distinct(Span.service_name)))) or 0
            error_count = session.scalar(select(func.count()).select_from(Span).where(Span.status_code == "ERROR")) or 0
            return OtelStats(
                traces=int(trace_count),
                spans=int(span_count),
                logs=int(log_count),
                metrics=int(metric_count),
                services=int(service_count),
                error_spans=int(error_count),
            )

    @staticmethod
    def _trace_summary(row: Trace) -> dict[str, object]:
        return {
            "trace_id": row.trace_id,
            "root_service_name": row.root_service_name,
            "root_span_name": row.root_span_name,
            "status_code": row.status_code,
            "span_count": row.span_count,
            "duration_ms": round(row.duration_ms, 3),
            "start_time": row.start_time.isoformat() if row.start_time else None,
            "end_time": row.end_time.isoformat() if row.end_time else None,
        }

    @staticmethod
    def _log_summary(row: LogEntry) -> dict[str, object]:
        return {
            "id": row.id,
            "timestamp": row.timestamp.isoformat() if row.timestamp else None,
            "observed_time": row.observed_time.isoformat() if row.observed_time else None,
            "service_name": row.service_name,
            "severity_text": row.severity_text,
            "body": row.body,
            "trace_id": row.trace_id,
            "span_id": row.span_id,
            "attributes": json.loads(row.attributes_json or "{}"),
        }

    @staticmethod
    def _metric_summary(row: MetricPoint) -> dict[str, object]:
        return {
            "id": row.id,
            "service_name": row.service_name,
            "name": row.metric_name,
            "metric_type": row.metric_type,
            "unit": row.unit,
            "temporality": row.temporality,
            "start_time": row.start_time.isoformat() if row.start_time else None,
            "time": row.time.isoformat() if row.time else None,
            "value": row.value,
            "attributes": json.loads(row.attributes_json or "{}"),
        }

    @staticmethod
    def _trace_timing(spans: list[SpanRecord]) -> tuple[datetime | None, datetime | None, float]:
        parsed_starts = [_parse_timestamp(span.start_time) for span in spans if span.start_time]
        parsed_ends = [_parse_timestamp(span.end_time) for span in spans if span.end_time]
        start_time = min((value for value in parsed_starts if value is not None), default=None)
        end_time = max((value for value in parsed_ends if value is not None), default=None)
        duration_ms = 0.0
        if start_time is not None and end_time is not None:
            duration_ms = max(0.0, (end_time - start_time).total_seconds() * 1000.0)
        return start_time, end_time, duration_ms

    # ------------------------------------------------------------------
    # Advanced query methods
    # ------------------------------------------------------------------

    def get_traces_in_range(
        self,
        start: datetime | None = None,
        end: datetime | None = None,
        service: str | None = None,
        status: str | None = None,
        limit: int = 100,
    ) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            stmt = select(Trace).order_by(Trace.start_time.desc().nullslast(), Trace.id.desc()).limit(limit)
            if start is not None:
                stmt = stmt.where(Trace.start_time >= start)
            if end is not None:
                stmt = stmt.where(Trace.start_time <= end)
            if service:
                stmt = stmt.where(Trace.root_service_name == service)
            if status:
                stmt = stmt.where(Trace.status_code == status.upper())
            return [self._trace_summary(row) for row in session.scalars(stmt)]

    def get_logs_in_range(
        self,
        start: datetime | None = None,
        end: datetime | None = None,
        service: str | None = None,
        severity: str | None = None,
        limit: int = 100,
    ) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            stmt = select(LogEntry).order_by(LogEntry.timestamp.desc().nullslast(), LogEntry.id.desc()).limit(limit)
            if start is not None:
                stmt = stmt.where(LogEntry.timestamp >= start)
            if end is not None:
                stmt = stmt.where(LogEntry.timestamp <= end)
            if service:
                stmt = stmt.where(LogEntry.service_name == service)
            if severity:
                stmt = stmt.where(LogEntry.severity_text == severity.upper())
            return [self._log_summary(row) for row in session.scalars(stmt)]

    def get_metrics_in_range(
        self,
        start: datetime | None = None,
        end: datetime | None = None,
        service: str | None = None,
        metric_name: str | None = None,
        limit: int = 100,
    ) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            stmt = select(MetricPoint).order_by(MetricPoint.time.desc().nullslast(), MetricPoint.id.desc()).limit(limit)
            if start is not None:
                stmt = stmt.where(MetricPoint.time >= start)
            if end is not None:
                stmt = stmt.where(MetricPoint.time <= end)
            if service:
                stmt = stmt.where(MetricPoint.service_name == service)
            if metric_name:
                stmt = stmt.where(MetricPoint.metric_name == metric_name)
            return [self._metric_summary(row) for row in session.scalars(stmt)]

    def search_logs(self, query: str, limit: int = 50) -> list[dict[str, object]]:
        """Full-text scan on log body (case-insensitive substring match)."""
        with session_scope(self.session_factory) as session:
            stmt = (
                select(LogEntry)
                .where(LogEntry.body.ilike(f"%{query}%"))
                .order_by(LogEntry.timestamp.desc().nullslast(), LogEntry.id.desc())
                .limit(limit)
            )
            return [self._log_summary(row) for row in session.scalars(stmt)]

    def get_dashboard_aggregation(self) -> dict[str, object]:
        """Return aggregated stats suitable for a dashboard overview."""
        stats = self.get_stats()
        services = self.list_services()
        return {
            "totals": stats.to_dict(),
            "services": [s.to_dict() for s in services],
            "error_rate": round(stats.error_spans / max(stats.spans, 1), 4),
        }

    def vacuum(self) -> None:
        """Run SQLite VACUUM to reclaim disk space."""
        with self.engine.connect() as conn:
            conn.execute(func.text("VACUUM"))

    def delete_signals_before(self, signal: str, cutoff: datetime) -> int:
        """Delete rows older than *cutoff* for the given signal type.

        Returns number of deleted rows.
        """
        from sqlalchemy import delete as _delete

        model_map = {
            "traces": Trace,
            "logs": LogEntry,
            "metrics": MetricPoint,
        }
        model = model_map.get(signal)
        if model is None:
            return 0

        time_col = {
            Trace: Trace.start_time,
            LogEntry: LogEntry.timestamp,
            MetricPoint: MetricPoint.time,
        }[model]

        with session_scope(self.session_factory) as session:
            result = session.execute(
                _delete(model).where(time_col < cutoff).returning(model.id)
            )
            return result.rowcount or 0

    def fetch_signals_before(self, signal: str, cutoff: datetime) -> list[dict]:
        """Fetch records older than *cutoff* for archiving."""
        if signal == "traces":
            return self.get_traces_in_range(end=cutoff, limit=10_000)
        if signal == "logs":
            return self.get_logs_in_range(end=cutoff, limit=10_000)
        if signal == "metrics":
            return self.get_metrics_in_range(end=cutoff, limit=10_000)
        return []
