from __future__ import annotations

import json
from datetime import timezone
from typing import Optional

import asyncpg
import structlog

from .envelope import Envelope

_log = structlog.get_logger(__name__)


async def insert_envelope(conn: asyncpg.Connection, env: Envelope) -> None:
    """Insert envelope into kb-postgres. Silently ignores duplicate ids (ON CONFLICT DO NOTHING)."""
    geo_param = json.dumps(env.geo) if env.geo is not None else None
    await conn.execute(
        """
        INSERT INTO envelope (id, source, ts, geo, raw_ref, entities)
        VALUES ($1, $2, $3, $4::jsonb, $5, $6::jsonb)
        ON CONFLICT (id) DO NOTHING
        """,
        env.id,
        env.source,
        env.ts,
        geo_param,
        env.raw_ref,
        json.dumps(env.entities),
    )
    _log.info("envelope.inserted", id=env.id, source=env.source)


async def get_envelope(conn: asyncpg.Connection, envelope_id: str) -> Optional[Envelope]:
    """Fetch a single envelope by id. Returns None if not found."""
    row = await conn.fetchrow(
        "SELECT id, source, ts, geo, raw_ref, entities FROM envelope WHERE id = $1",
        envelope_id,
    )
    if row is None:
        return None

    ts = row["ts"]
    if ts.tzinfo is None:
        ts = ts.replace(tzinfo=timezone.utc)

    # asyncpg decodes JSONB columns to Python objects; handle both str and object
    geo = _decode_jsonb(row["geo"])
    entities = _decode_jsonb(row["entities"]) or []

    return Envelope(
        id=row["id"],
        source=row["source"],
        ts=ts,
        geo=geo,
        raw_ref=row["raw_ref"],
        entities=entities,
    )


def _decode_jsonb(value: object) -> object:
    """Decode a JSONB value that may arrive as a string or already-decoded object."""
    if value is None:
        return None
    if isinstance(value, str):
        return json.loads(value)
    return value
