"""Log entry types for Lattice."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

import deal

from lattice.core.types.enums import Role


@dataclass(frozen=True)
class LogEntry:
    """A single log entry in store.db.

    >>> entry = LogEntry(
    ...     external_id="uuid-123",
    ...     session_id="session-abc",
    ...     timestamp="2026-02-17T10:00:00Z",
    ...     role=Role.USER,
    ...     content="Hello",
    ... )
    >>> entry.role
    <Role.USER: 'user'>
    """

    external_id: str
    session_id: str
    timestamp: str
    role: Role
    content: str
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class Session:
    """A conversation session with its logs.

    >>> session = Session(
    ...     session_id="session-abc",
    ...     logs=[
    ...         LogEntry(
    ...             external_id="uuid-1",
    ...             session_id="session-abc",
    ...             timestamp="2026-02-17T10:00:00Z",
    ...             role=Role.USER,
    ...             content="Hello",
    ...         ),
    ...     ],
    ... )
    >>> len(session.logs)
    1
    """

    session_id: str
    logs: list[LogEntry] = field(default_factory=list)
    started_at: str | None = None
    ended_at: str | None = None


@deal.pre(lambda logs: isinstance(logs, list))
@deal.pre(lambda logs: all(isinstance(log, LogEntry) for log in logs))
@deal.post(lambda result: isinstance(result, list))
@deal.post(lambda result: all(isinstance(s, Session) for s in result))
def group_logs_into_sessions(logs: list[LogEntry]) -> list[Session]:
    """Group a flat list of LogEntry into Session objects.

    Groups by session_id, orders entries by timestamp within each session.

    Args:
        logs: Flat list of LogEntry objects.

    Returns:
        List of Session objects, each containing its logs in timestamp order.

    >>> logs = [
    ...     LogEntry("id1", "session-a", "2026-01-01T10:00:00Z", Role.USER, "Hi"),
    ...     LogEntry("id2", "session-a", "2026-01-01T10:01:00Z", Role.ASSISTANT, "Hello"),
    ...     LogEntry("id3", "session-b", "2026-01-01T11:00:00Z", Role.USER, "Bye"),
    ... ]
    >>> sessions = group_logs_into_sessions(logs)
    >>> len(sessions)
    2
    >>> sessions[0].session_id
    'session-a'
    >>> len(sessions[0].logs)
    2
    """
    if not logs:
        return []

    # Group by session_id
    sessions_dict: dict[str, list[LogEntry]] = {}
    for log in logs:
        if log.session_id not in sessions_dict:
            sessions_dict[log.session_id] = []
        sessions_dict[log.session_id].append(log)

    # Build Session objects with sorted logs
    sessions: list[Session] = []
    for session_id, session_logs in sessions_dict.items():
        sorted_logs = sorted(session_logs, key=lambda x: x.timestamp)
        sessions.append(
            Session(
                session_id=session_id,
                logs=sorted_logs,
                started_at=sorted_logs[0].timestamp if sorted_logs else None,
                ended_at=sorted_logs[-1].timestamp if sorted_logs else None,
            )
        )

    # Sort sessions by started_at
    return sorted(sessions, key=lambda s: s.started_at or "")
