"""Temporal Context API: chronological knowledge retrieval.

Queries NATS events and ChromaDB with a time dimension to answer:
"What do we know about X, chronologically?"
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any

logger = logging.getLogger(__name__)

# Default config from environment
NATS_URL = os.environ.get("NATS_URL", "nats://localhost:4222")
CHROMADB_URL = os.environ.get("CHROMADB_URL", "http://localhost:8000")


@dataclass
class TemporalEntry:
    """A knowledge entry with temporal metadata."""

    timestamp: datetime
    source: str  # "nats" or "chromadb"
    topic: str
    content: str
    metadata: dict[str, Any] = field(default_factory=dict)
    relevance_score: float = 0.0


@dataclass
class TemporalQuery:
    """Query parameters for temporal context retrieval."""

    topic: str
    start_time: datetime | None = None
    end_time: datetime | None = None
    limit: int = 50
    sources: list[str] = field(default_factory=lambda: ["nats", "chromadb"])


class TemporalContext:
    """Retrieves chronological knowledge from NATS events and ChromaDB.

    Usage:
        ctx = TemporalContext()
        entries = await ctx.query(TemporalQuery(topic="ssl-cert"))
    """

    def __init__(
        self,
        nats_url: str | None = None,
        chromadb_url: str | None = None,
    ) -> None:
        self.nats_url = nats_url or NATS_URL
        self.chromadb_url = chromadb_url or CHROMADB_URL
        self._nats_client: Any = None
        self._chroma_client: Any = None

    async def connect(self) -> None:
        """Initialize connections to NATS and ChromaDB."""
        try:
            import nats
            self._nats_client = await nats.connect(self.nats_url)
            logger.info("Connected to NATS: %s", self.nats_url)
        except Exception:
            logger.exception("Failed to connect to NATS")

        try:
            import chromadb
            self._chroma_client = chromadb.HttpClient(host=self.chromadb_url)
            logger.info("Connected to ChromaDB: %s", self.chromadb_url)
        except Exception:
            logger.exception("Failed to connect to ChromaDB")

    async def query(self, query: TemporalQuery) -> list[TemporalEntry]:
        """Query temporal context across configured sources.

        Returns entries sorted chronologically (oldest first).
        """
        entries: list[TemporalEntry] = []

        if "nats" in query.sources and self._nats_client:
            nats_entries = await self._query_nats(query)
            entries.extend(nats_entries)

        if "chromadb" in query.sources and self._chroma_client:
            chroma_entries = self._query_chromadb(query)
            entries.extend(chroma_entries)

        # Sort chronologically
        entries.sort(key=lambda e: e.timestamp)

        # Apply limit
        if query.limit:
            entries = entries[:query.limit]

        return entries

    async def _query_nats(self, query: TemporalQuery) -> list[TemporalEntry]:
        """Query NATS JetStream for historical events matching the topic."""
        entries: list[TemporalEntry] = []
        try:
            js = self._nats_client.jetstream()
            subject = f"darkplex.*.{query.topic}.>"

            # Get messages from the stream
            sub = await js.subscribe(subject, ordered_consumer=True)
            count = 0
            async for msg in sub.messages:
                if count >= query.limit:
                    break

                timestamp = datetime.fromtimestamp(
                    msg.headers.get("Nats-Time-Stamp", 0) if msg.headers else 0,
                    tz=timezone.utc,
                )

                if query.start_time and timestamp < query.start_time:
                    continue
                if query.end_time and timestamp > query.end_time:
                    continue

                entries.append(TemporalEntry(
                    timestamp=timestamp,
                    source="nats",
                    topic=query.topic,
                    content=msg.data.decode() if msg.data else "",
                    metadata={"subject": msg.subject},
                ))
                count += 1

        except Exception:
            logger.exception("NATS temporal query failed for topic: %s", query.topic)

        return entries

    def _query_chromadb(self, query: TemporalQuery) -> list[TemporalEntry]:
        """Query ChromaDB for semantically relevant entries with time filtering."""
        entries: list[TemporalEntry] = []
        try:
            collection = self._chroma_client.get_or_create_collection("darkplex_knowledge")

            where_filter: dict[str, Any] = {}
            if query.start_time:
                where_filter["timestamp"] = {"$gte": query.start_time.isoformat()}
            if query.end_time:
                if "timestamp" in where_filter:
                    where_filter = {
                        "$and": [
                            {"timestamp": {"$gte": query.start_time.isoformat()}},
                            {"timestamp": {"$lte": query.end_time.isoformat()}},
                        ]
                    }
                else:
                    where_filter["timestamp"] = {"$lte": query.end_time.isoformat()}

            results = collection.query(
                query_texts=[query.topic],
                n_results=query.limit,
                where=where_filter if where_filter else None,
            )

            if results and results.get("documents"):
                for i, doc in enumerate(results["documents"][0]):
                    meta = results["metadatas"][0][i] if results.get("metadatas") else {}
                    ts_str = meta.get("timestamp", "")
                    try:
                        ts = datetime.fromisoformat(ts_str)
                    except (ValueError, TypeError):
                        ts = datetime.now(timezone.utc)

                    entries.append(TemporalEntry(
                        timestamp=ts,
                        source="chromadb",
                        topic=query.topic,
                        content=doc,
                        metadata=meta,
                        relevance_score=results["distances"][0][i] if results.get("distances") else 0.0,
                    ))

        except Exception:
            logger.exception("ChromaDB temporal query failed for topic: %s", query.topic)

        return entries

    async def close(self) -> None:
        """Close connections."""
        if self._nats_client:
            await self._nats_client.close()