planet/backend/app/services/vessel_ais_aggregation.py

"""AIS raw observation and aggregation support for vessel collectors."""

from datetime import UTC, datetime, timedelta
from hashlib import sha256
import json
from typing import Any, Iterable

from sqlalchemy import select
from sqlalchemy import Float
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.vessel import AISConflictRecord, AISRawObservation, AISSourceHealth
from app.services.vessel_aggregation_strategy import (
    DEFAULT_STRATEGY,
    load_strategy,
)
from app.services.vessel_types import normalize_vessel_type_name

VESSEL_AIS_SCHEMA = "vessel_ais"
DEFAULT_AGGREGATION_WINDOW_HOURS = 24
DEFAULT_SNAPSHOT_WINDOW_MINUTES = 60
MAX_SNAPSHOT_LIMIT = 5000
MAX_SNAPSHOT_CANDIDATE_MULTIPLIER = 20
MAX_SNAPSHOT_CANDIDATE_OBSERVATIONS = 100_000
BARENTSWATCH_DELIVERY_MODE = "polling"
BARENTSWATCH_TRANSPORT = "http"
AISSTREAM_DELIVERY_MODE = "realtime_stream"
AISSTREAM_TRANSPORT = "websocket"
DELIVERY_MODE_PRIORITY = {
    "realtime_stream": 40,
    "batch_stream": 30,
    "polling": 20,
    "snapshot": 10,
}
DYNAMIC_FIELDS = ("lat", "lon", "sog", "cog", "heading", "nav_status")
CONFLICT_FIELDS = (
    "name",
    "callsign",
    "imo",
    "flag",
    "vessel_type",
    "vessel_type_name",
    "length",
    "width",
    "draught",
)


def _json_default(value: Any) -> Any:
    if isinstance(value, datetime):
        return value.astimezone(UTC).isoformat()
    return str(value)


def _stable_payload(value: Any) -> str:
    return json.dumps(value, sort_keys=True, separators=(",", ":"), default=_json_default)


def _jsonable(value: Any) -> Any:
    if isinstance(value, datetime):
        return value.astimezone(UTC).isoformat()
    if isinstance(value, dict):
        return {str(key): _jsonable(item) for key, item in value.items()}
    if isinstance(value, list):
        return [_jsonable(item) for item in value]
    return value


def _coerce_datetime(value: Any) -> datetime | None:
    if isinstance(value, datetime):
        return value if value.tzinfo else value.replace(tzinfo=UTC)
    if isinstance(value, (int, float)):
        timestamp = float(value)
        if timestamp > 10_000_000_000:
            timestamp /= 1000
        return datetime.fromtimestamp(timestamp, UTC)
    if isinstance(value, str) and value:
        try:
            parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
            return parsed if parsed.tzinfo else parsed.replace(tzinfo=UTC)
        except ValueError:
            return None
    return None


def build_observation_hash(
    *,
    source: str,
    entity_key: str,
    message_type: str | None,
    observed_at: datetime,
    normalized_payload: dict[str, Any],
    source_message_id: str | None = None,
) -> str:
    """Build a deterministic idempotency key for one source-level AIS observation."""

    if source_message_id:
        basis = {
            "source": source,
            "entity_key": entity_key,
            "source_message_id": source_message_id,
        }
    else:
        basis = {
            "source": source,
            "entity_key": entity_key,
            "message_type": message_type,
            "observed_at": observed_at.astimezone(UTC).isoformat(),
            "payload": normalized_payload,
        }
    return sha256(_stable_payload(basis).encode("utf-8")).hexdigest()


def build_field_conflict_candidates(
    observations: Iterable[AISRawObservation],
    fields: Iterable[str] = CONFLICT_FIELDS,
) -> list[dict[str, Any]]:
    """Return current field disagreements from raw observations without mutating state."""

    candidates_by_field: dict[str, dict[str, Any]] = {}
    for observation in observations:
        payload = observation.normalized_payload or {}
        for field in fields:
            value = payload.get(field)
            if value in (None, ""):
                continue
            field_candidates = candidates_by_field.setdefault(field, {})
            field_candidates[observation.source] = value

    conflicts = []
    for field, candidates in sorted(candidates_by_field.items()):
        unique_values = {_stable_payload(value) for value in candidates.values()}
        if len(unique_values) <= 1:
            continue
        conflicts.append(
            {
                "field": field,
                "candidates": candidates,
                "status": "candidate",
            }
        )
    return conflicts


def _payload_value(payload: dict[str, Any], field: str) -> Any:
    value = payload.get(field)
    return None if value in (None, "") else value


def _clean_text(value: Any) -> str | None:
    if value in (None, ""):
        return None
    text = str(value).strip()
    return text or None


def _raw_metadata_value(observation: AISRawObservation, field: str) -> Any:
    raw_payload = observation.raw_payload or {}
    metadata = raw_payload.get("MetaData") if isinstance(raw_payload, dict) else None
    if not isinstance(metadata, dict):
        return None
    if field == "name":
        return _clean_text(metadata.get("ShipName") or metadata.get("ship_name") or metadata.get("name"))
    return None


def _delivery_priority(observation: AISRawObservation) -> int:
    return DELIVERY_MODE_PRIORITY.get(str(observation.delivery_mode or ""), 0)


def _has_valid_position(payload: dict[str, Any]) -> bool:
    try:
        lat = float(payload.get("lat"))
        lon = float(payload.get("lon"))
    except (TypeError, ValueError):
        return False
    return -90 <= lat <= 90 and -180 <= lon <= 180


def _is_future_observation(observation: AISRawObservation, now: datetime) -> bool:
    return observation.observed_at > now


def _strategy_source_rank(
    source: str,
    strategy: dict[str, Any],
) -> int:
    priority = (strategy.get("vessel_ais") or {}).get("source_priority") or []
    if source in priority:
        return len(priority) - priority.index(source)
    return 0


def _is_stream_stale(
    observation: AISRawObservation,
    *,
    now: datetime,
    strategy: dict[str, Any],
) -> bool:
    delivery_mode = str(observation.delivery_mode or "")
    freshness = (strategy.get("vessel_ais") or {}).get("freshness") or {}
    if delivery_mode == "realtime_stream":
        window = int(freshness.get("realtime_stream_seconds", 0) or 0)
    else:
        window = int(freshness.get("polling_seconds", 0) or 0)
    if window <= 0:
        return False
    return (now - observation.observed_at).total_seconds() > window


def _select_position_observation(
    observations: list[AISRawObservation],
    *,
    now: datetime,
    strategy: dict[str, Any] | None = None,
) -> tuple[AISRawObservation | None, list[str]]:
    strategy = strategy or DEFAULT_STRATEGY
    rejected_flags: list[str] = []
    fresh_candidates: list[AISRawObservation] = []
    stale_candidates: list[AISRawObservation] = []
    for observation in observations:
        payload = observation.normalized_payload or {}
        if not _has_valid_position(payload):
            rejected_flags.append("invalid_position")
            continue
        if _is_future_observation(observation, now):
            rejected_flags.append("future_timestamp")
            continue
        if _is_stream_stale(observation, now=now, strategy=strategy):
            stale_candidates.append(observation)
            rejected_flags.append("freshness_fallback")
            continue
        fresh_candidates.append(observation)

    candidates = fresh_candidates or stale_candidates
    if not candidates:
        return None, sorted(set(rejected_flags))

    candidates.sort(
        key=lambda item: (
            item.observed_at,
            _delivery_priority(item),
            _strategy_source_rank(item.source, strategy),
            item.collected_at,
            item.id or 0,
        ),
        reverse=True,
    )
    return candidates[0], sorted(set(rejected_flags))


def _select_static_field(
    observations: list[AISRawObservation],
    field: str,
    strategy: dict[str, Any] | None = None,
) -> tuple[Any, str | None, str | None]:
    strategy = strategy or DEFAULT_STRATEGY
    candidates = []
    for observation in observations:
        value = _payload_value(observation.normalized_payload or {}, field)
        if value is None:
            value = _raw_metadata_value(observation, field)
        if value is None:
            continue
        candidates.append((observation, value))

    if not candidates:
        return None, None, None

    field_rules = (strategy.get("vessel_ais") or {}).get("field_rules") or {}
    rule = field_rules.get(field) or {"mode": "source_priority"}
    mode = rule.get("mode")

    if mode == "locked":
        locked_source = rule.get("locked_source")
        for observation, value in candidates:
            if observation.source == locked_source:
                return value, observation.source, "locked"

    if mode in ("source_priority", "locked"):
        priority = rule.get("source_priority") or (strategy.get("vessel_ais") or {}).get("source_priority") or []
        ranked = sorted(
            candidates,
            key=lambda item: (
                priority.index(item[0].source) if item[0].source in priority else len(priority) + 1,
                -_delivery_priority(item[0]),
                -(item[0].observed_at.timestamp() if item[0].observed_at else 0),
            ),
        )
        observation, value = ranked[0]
        return value, observation.source, "source_priority"

    if mode == "newest":
        ranked = sorted(
            candidates,
            key=lambda item: (item[0].observed_at, _delivery_priority(item[0]), item[0].id or 0),
            reverse=True,
        )
        observation, value = ranked[0]
        return value, observation.source, "newest_observation"

    # default / non_empty: prefer delivery mode priority, then newest
    candidates.sort(
        key=lambda item: (
            _delivery_priority(item[0]),
            item[0].observed_at,
            item[0].collected_at,
            item[0].id or 0,
        ),
        reverse=True,
    )
    selected_observation, selected_value = candidates[0]
    unique_values = {_stable_payload(value) for _, value in candidates}
    reason = "delivery_mode_priority" if len(unique_values) > 1 else "non_empty_priority"
    return selected_value, selected_observation.source, reason


def _build_source_summary(observations: list[AISRawObservation]) -> dict[str, dict[str, Any]]:
    summary: dict[str, dict[str, Any]] = {}
    for observation in observations:
        source_summary = summary.setdefault(
            observation.source,
            {
                "observation_count": 0,
                "latest_observed_at": None,
                "delivery_mode": observation.delivery_mode,
                "transport": observation.transport,
                "message_types": [],
            },
        )
        source_summary["observation_count"] += 1
        latest_observed_at = source_summary["latest_observed_at"]
        if latest_observed_at is None or observation.observed_at > latest_observed_at:
            source_summary["latest_observed_at"] = observation.observed_at
        if observation.message_type and observation.message_type not in source_summary["message_types"]:
            source_summary["message_types"].append(observation.message_type)
    return summary


def _build_aggregated_vessel(
    entity_key: str,
    observations: list[AISRawObservation],
    *,
    now: datetime,
    strategy: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
    strategy = strategy or DEFAULT_STRATEGY
    position_observation, rejected_flags = _select_position_observation(
        observations, now=now, strategy=strategy
    )
    if position_observation is None:
        return None

    payload = position_observation.normalized_payload or {}
    mmsi = int(entity_key)
    result: dict[str, Any] = {
        "mmsi": mmsi,
        "lat": float(payload["lat"]),
        "lon": float(payload["lon"]),
        "received_at": position_observation.observed_at,
        "field_sources": {},
        "selected_reasons": {},
        "source_summary": _build_source_summary(observations),
        "quality_flags": sorted(
            set((position_observation.quality_flags or []) + rejected_flags)
        ),
        "aggregation_strategy_version": int(strategy.get("version") or 0),
    }

    for field in DYNAMIC_FIELDS:
        value = _payload_value(payload, field)
        if field in ("lat", "lon") or value is not None:
            result[field] = value
            result["field_sources"][field] = position_observation.source
            result["selected_reasons"][field] = "newest_observation"

    for field in CONFLICT_FIELDS:
        selected_value, selected_source, reason = _select_static_field(
            observations, field, strategy=strategy
        )
        if selected_value is None:
            continue
        result[field] = selected_value
        result["field_sources"][field] = selected_source
        result["selected_reasons"][field] = reason

    result["name"] = result.get("name") or f"MMSI {mmsi}"
    result["vessel_type_name"] = result.get("vessel_type_name") or normalize_vessel_type_name(
        result.get("vessel_type")
    )
    return result


async def _upsert_conflict_records(
    db: AsyncSession,
    entity_key: str,
    observations: list[AISRawObservation],
    aggregated: dict[str, Any],
) -> int:
    conflicts = build_field_conflict_candidates(observations)
    now = datetime.now(UTC)
    for conflict in conflicts:
        field = conflict["field"]
        result = await db.execute(
            select(AISConflictRecord)
            .where(AISConflictRecord.target_schema == VESSEL_AIS_SCHEMA)
            .where(AISConflictRecord.entity_key == entity_key)
            .where(AISConflictRecord.field == field)
            .limit(1)
        )
        record = result.scalar_one_or_none()
        if record is None:
            record = AISConflictRecord(
                target_schema=VESSEL_AIS_SCHEMA,
                entity_key=entity_key,
                field=field,
            )
            db.add(record)
        record.candidates = conflict["candidates"]
        record.selected_source = (aggregated.get("field_sources") or {}).get(field)
        record.selected_value = aggregated.get(field)
        record.selected_reason = (aggregated.get("selected_reasons") or {}).get(field)
        record.resolved_by = "system"
        record.status = "open"
        record.updated_at = now
    return len(conflicts)


def _group_observations(observations: Iterable[AISRawObservation]) -> dict[str, list[AISRawObservation]]:
    grouped: dict[str, list[AISRawObservation]] = {}
    for observation in observations:
        grouped.setdefault(str(observation.entity_key), []).append(observation)
    return grouped


async def record_vessel_ais_observation(
    db: AsyncSession,
    *,
    source: str,
    normalized_payload: dict[str, Any],
    raw_payload: dict[str, Any] | None = None,
    delivery_mode: str,
    transport: str,
    message_type: str | None = "PositionReport",
    source_message_id: str | None = None,
    observed_at: datetime | None = None,
    collected_at: datetime | None = None,
    quality_flags: list[str] | None = None,
) -> AISRawObservation | None:
    """Insert one raw observation if the source-level fact has not already been stored."""

    entity_key = str(normalized_payload["mmsi"])
    collected_at = collected_at or datetime.now(UTC)
    observed_at = (
        _coerce_datetime(observed_at)
        or _coerce_datetime(normalized_payload.get("received_at"))
        or collected_at
    )
    normalized_json = _jsonable(normalized_payload)
    raw_json = _jsonable(raw_payload or {})

    observation_hash = build_observation_hash(
        source=source,
        entity_key=entity_key,
        message_type=message_type,
        observed_at=observed_at,
        normalized_payload=normalized_json,
        source_message_id=source_message_id,
    )
    existing_result = await db.execute(
        select(AISRawObservation.id).where(AISRawObservation.observation_hash == observation_hash)
    )
    if existing_result.scalar_one_or_none() is not None:
        return None

    observation = AISRawObservation(
        target_schema=VESSEL_AIS_SCHEMA,
        source=source,
        entity_key=entity_key,
        delivery_mode=delivery_mode,
        transport=transport,
        message_type=message_type,
        source_message_id=source_message_id,
        observation_hash=observation_hash,
        observed_at=observed_at,
        collected_at=collected_at,
        normalized_payload=normalized_json,
        raw_payload=raw_json,
        quality_flags=quality_flags or [],
    )
    db.add(observation)
    return observation


async def aggregate_vessel_observations(
    db: AsyncSession,
    observations: Iterable[AISRawObservation],
    *,
    write_conflicts: bool = False,
    strategy: dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    strategy = strategy if strategy is not None else await _safe_load_strategy(db)
    now = datetime.now(UTC)
    vessels = []
    for entity_key, entity_observations in _group_observations(observations).items():
        aggregated = _build_aggregated_vessel(
            entity_key, entity_observations, now=now, strategy=strategy
        )
        if aggregated is None:
            continue
        if write_conflicts:
            aggregated["conflict_count"] = await _upsert_conflict_records(
                db,
                entity_key,
                entity_observations,
                aggregated,
            )
        else:
            aggregated["conflict_count"] = len(build_field_conflict_candidates(entity_observations))
        vessels.append(aggregated)

    vessels.sort(key=lambda item: item.get("received_at") or datetime.min.replace(tzinfo=UTC), reverse=True)
    return vessels


async def _safe_load_strategy(db: AsyncSession) -> dict[str, Any]:
    """Tolerate fake test sessions where load_strategy may misbehave."""
    try:
        return await load_strategy(db)
    except Exception:
        return DEFAULT_STRATEGY


async def get_aggregated_vessels(
    db: AsyncSession,
    *,
    bbox: tuple[float, float, float, float] | None = None,
    limit: int | None = None,
    observed_since: datetime | None = None,
) -> list[dict[str, Any]]:
    observed_since = observed_since or (
        datetime.now(UTC) - timedelta(hours=DEFAULT_AGGREGATION_WINDOW_HOURS)
    )
    stmt = (
        select(AISRawObservation)
        .where(AISRawObservation.target_schema == VESSEL_AIS_SCHEMA)
        .where(AISRawObservation.observed_at >= observed_since)
        .order_by(AISRawObservation.observed_at.desc(), AISRawObservation.id.desc())
    )
    if limit and limit > 0:
        stmt = stmt.limit(max(limit * 20, limit))

    result = await db.execute(stmt)
    if not hasattr(result, "scalars"):
        return []
    vessels = await aggregate_vessel_observations(db, result.scalars().all())

    if bbox is not None:
        lon_min, lat_min, lon_max, lat_max = bbox
        vessels = [
            vessel
            for vessel in vessels
            if lon_min <= float(vessel["lon"]) <= lon_max
            and lat_min <= float(vessel["lat"]) <= lat_max
        ]

    if limit and limit > 0:
        return vessels[:limit]
    return vessels


async def get_aggregated_vessels_snapshot(
    db: AsyncSession,
    *,
    bbox: tuple[float, float, float, float],
    limit: int = 1000,
    observed_since: datetime | None = None,
) -> list[dict[str, Any]]:
    """Return a bounded viewport snapshot without loading the global AIS window."""

    observed_since = observed_since or (
        datetime.now(UTC) - timedelta(minutes=DEFAULT_SNAPSHOT_WINDOW_MINUTES)
    )
    safe_limit = min(max(int(limit or 1000), 1), MAX_SNAPSHOT_LIMIT)
    candidate_limit = min(
        max(safe_limit * MAX_SNAPSHOT_CANDIDATE_MULTIPLIER, safe_limit),
        MAX_SNAPSHOT_CANDIDATE_OBSERVATIONS,
    )
    lon_min, lat_min, lon_max, lat_max = bbox
    payload_lon = AISRawObservation.normalized_payload["lon"].as_string().cast(Float)
    payload_lat = AISRawObservation.normalized_payload["lat"].as_string().cast(Float)

    stmt = (
        select(AISRawObservation)
        .where(AISRawObservation.target_schema == VESSEL_AIS_SCHEMA)
        .where(AISRawObservation.observed_at >= observed_since)
        .where(payload_lon >= lon_min)
        .where(payload_lon <= lon_max)
        .where(payload_lat >= lat_min)
        .where(payload_lat <= lat_max)
        .order_by(AISRawObservation.observed_at.desc(), AISRawObservation.id.desc())
        .limit(candidate_limit)
    )

    result = await db.execute(stmt)
    if not hasattr(result, "scalars"):
        return []
    vessels = await aggregate_vessel_observations(db, result.scalars().all())
    return vessels[:safe_limit]


async def get_aggregated_vessel(db: AsyncSession, mmsi: int) -> dict[str, Any] | None:
    observations = await get_vessel_raw_observations(db, mmsi, limit=1000)
    vessels = await aggregate_vessel_observations(db, observations)
    return vessels[0] if vessels else None


async def get_aggregated_vessel_track(
    db: AsyncSession,
    mmsi: int,
    *,
    cutoff: datetime,
) -> list[dict[str, Any]]:
    result = await db.execute(
        select(AISRawObservation)
        .where(AISRawObservation.target_schema == VESSEL_AIS_SCHEMA)
        .where(AISRawObservation.entity_key == str(mmsi))
        .where(AISRawObservation.observed_at >= cutoff)
        .order_by(AISRawObservation.observed_at.asc(), AISRawObservation.id.asc())
    )
    if not hasattr(result, "scalars"):
        return []

    points: list[dict[str, Any]] = []
    seen: set[tuple[str, float, float, str]] = set()
    for observation in result.scalars().all():
        payload = observation.normalized_payload or {}
        if not _has_valid_position(payload):
            continue
        lat = float(payload["lat"])
        lon = float(payload["lon"])
        key = (
            observation.observed_at.isoformat(),
            round(lat, 5),
            round(lon, 5),
            observation.source,
        )
        if key in seen:
            continue
        seen.add(key)
        points.append(
            {
                "lat": lat,
                "lon": lon,
                "observed_at": observation.observed_at,
                "source": observation.source,
                "selected_reason": "track_timeline",
                "quality_flags": observation.quality_flags or [],
            }
        )
    return points


async def update_ais_source_health(
    db: AsyncSession,
    *,
    source: str,
    connection_state: str,
    observed_count: int = 0,
    last_seen_at: datetime | None = None,
    last_success_at: datetime | None = None,
    last_error: str | None = None,
    lag_seconds: float | None = None,
) -> AISSourceHealth:
    """Upsert the health row for an AIS source."""

    now = datetime.now(UTC)
    health = await db.get(AISSourceHealth, source)
    if health is None:
        health = AISSourceHealth(source=source)
        db.add(health)

    health.connection_state = connection_state
    health.last_seen_at = last_seen_at or health.last_seen_at
    health.last_success_at = last_success_at or health.last_success_at
    health.last_error = last_error
    health.message_rate = float(observed_count)
    health.lag_seconds = lag_seconds
    health.updated_at = now
    return health


async def count_unique_raw_vessel_mmsi(
    db: AsyncSession,
    *,
    observed_since: datetime | None = None,
) -> int:
    """Count unique raw vessel MMSI values for HUD counts; never aggregates."""
    from sqlalchemy import func as sa_func

    unique_mmsi_stmt = (
        select(AISRawObservation.entity_key)
        .where(AISRawObservation.target_schema == VESSEL_AIS_SCHEMA)
        .distinct()
    )
    if observed_since is not None:
        unique_mmsi_stmt = unique_mmsi_stmt.where(
            AISRawObservation.observed_at >= observed_since,
        )

    result = await db.execute(
        select(sa_func.count()).select_from(unique_mmsi_stmt.subquery()),
    )
    return int(result.scalar() or 0)


async def get_vessel_raw_observations(
    db: AsyncSession,
    mmsi: int,
    *,
    limit: int = 100,
) -> list[AISRawObservation]:
    result = await db.execute(
        select(AISRawObservation)
        .where(AISRawObservation.target_schema == VESSEL_AIS_SCHEMA)
        .where(AISRawObservation.entity_key == str(mmsi))
        .order_by(AISRawObservation.observed_at.desc(), AISRawObservation.id.desc())
        .limit(limit)
    )
    return list(result.scalars().all())


async def get_vessel_conflict_records(
    db: AsyncSession,
    mmsi: int,
) -> list[AISConflictRecord]:
    result = await db.execute(
        select(AISConflictRecord)
        .where(AISConflictRecord.target_schema == VESSEL_AIS_SCHEMA)
        .where(AISConflictRecord.entity_key == str(mmsi))
        .order_by(AISConflictRecord.updated_at.desc(), AISConflictRecord.id.desc())
    )
    return list(result.scalars().all())