"""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


@deal.inv(
    lambda self: (
        self.role == Role.TOOL
        or (
            self.tool_input is None
            and self.tool_status is None
            and self.tool_error is None
        )
    )
)
@deal.inv(lambda self: self.tool_status is None or self.tool_status != "")
@deal.inv(lambda self: self.tool_error is None or self.tool_error != "")
@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'>

    LogEntry with tool metadata fields (only valid for Role.TOOL):

    >>> tool_entry = LogEntry(
    ...     external_id="uuid-456",
    ...     session_id="session-abc",
    ...     timestamp="2026-02-17T10:01:00Z",
    ...     role=Role.TOOL,
    ...     content="Tool execution completed",
    ...     tool_input='{"query": "search term"}',
    ...     tool_status="success",
    ...     tool_error=None,
    ... )
    >>> tool_entry.tool_input
    '{"query": "search term"}'
    >>> tool_entry.tool_status
    'success'

    Backward compatibility - existing code works unchanged:

    >>> legacy = LogEntry("id1", "sess", "2026-01-01T00:00:00Z", Role.USER, "msg")
    >>> legacy.tool_input is None
    True
    >>> legacy.tool_status is None
    True

    Invariant validation - empty strings are rejected:

    >>> LogEntry("id", "sess", "t", Role.USER, "c", tool_status="")  # doctest: +IGNORE_EXCEPTION_DETAIL
    Traceback (most recent call last):
        ...
    InvContractError: ...

    Invariant validation - tool fields rejected for non-TOOL roles:

    >>> LogEntry("id", "sess", "t", Role.USER, "c", tool_input="data")  # doctest: +IGNORE_EXCEPTION_DETAIL
    Traceback (most recent call last):
        ...
    InvContractError: ...

    >>> LogEntry("id", "sess", "t", Role.ASSISTANT, "c", tool_status="success")  # doctest: +IGNORE_EXCEPTION_DETAIL
    Traceback (most recent call last):
        ...
    InvContractError: ...
    """

    external_id: str
    session_id: str
    timestamp: str
    role: Role
    content: str
    metadata: dict[str, Any] = field(default_factory=dict)
    tool_input: str | None = None
    tool_status: str | None = None
    tool_error: str | None = None

    # @invar:allow missing_contract: __post_init__ enforces @deal.inv invariants (frozen dataclass pattern)
    @deal.post(lambda result: result is None)
    def __post_init__(self) -> None:
        """Enforce invariants at construction time.

        deal.inv on frozen dataclasses documents invariants but doesn't enforce
        on construction. We check here to raise InvContractError early.
        """
        # Invariant: tool_* fields only valid for Role.TOOL
        if self.role != Role.TOOL:
            if self.tool_input is not None:
                raise deal.InvContractError(
                    f"tool_input must be None for role={self.role}, got {self.tool_input!r}"
                )
            if self.tool_status is not None:
                raise deal.InvContractError(
                    f"tool_status must be None for role={self.role}, got {self.tool_status!r}"
                )
            if self.tool_error is not None:
                raise deal.InvContractError(
                    f"tool_error must be None for role={self.role}, got {self.tool_error!r}"
                )
        # Invariant: tool_status and tool_error cannot be empty strings
        if self.tool_status == "":
            raise deal.InvContractError("tool_status cannot be empty string")
        if self.tool_error == "":
            raise deal.InvContractError("tool_error cannot be empty string")


@deal.inv(lambda self: len(self.session_id) > 0)
@deal.inv(lambda self: all(log.session_id == self.session_id for log in self.logs))
@deal.inv(lambda self: self.started_at is None or len(self.started_at) > 0)
@deal.inv(lambda self: self.ended_at is None or len(self.ended_at) > 0)
@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

    Invariant validation - empty session_id is rejected:

    >>> Session("", [])  # doctest: +IGNORE_EXCEPTION_DETAIL
    Traceback (most recent call last):
        ...
    InvContractError: ...

    Invariant validation - logs must match session_id:

    >>> Session(
    ...     session_id="session-a",
    ...     logs=[LogEntry("id", "session-b", "t", Role.USER, "c")],
    ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
    Traceback (most recent call last):
        ...
    InvContractError: ...

    Invariant validation - empty string timestamps rejected:

    >>> Session(session_id="s", logs=[], started_at="")  # doctest: +IGNORE_EXCEPTION_DETAIL
    Traceback (most recent call last):
        ...
    InvContractError: ...
    """

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


# @invar:allow missing_contract: Validation helper for precondition - cannot have own
# contracts since it must handle Hypothesis-generated invalid objects without raising.
def _log_entry_is_valid(log: LogEntry) -> bool:
    """Check if a LogEntry satisfies all its invariants.

    This is needed because deal's Hypothesis integration may construct
    LogEntry objects that bypass __post_init__ validation. The precondition
    must explicitly verify all LogEntry invariants.

    In normal code, all LogEntry objects are valid (construction rejects invalid).
    This function exists to catch Hypothesis-generated edge cases.

    >>> from lattice.core.types.enums import Role
    >>> # Valid USER entry
    >>> _log_entry_is_valid(LogEntry("id", "s", "t", Role.USER, "c"))
    True
    >>> # Valid TOOL entry with all tool fields
    >>> _log_entry_is_valid(LogEntry("id", "s", "t", Role.TOOL, "c", tool_input="x", tool_status="ok"))
    True
    >>> # Valid TOOL entry with minimal fields
    >>> _log_entry_is_valid(LogEntry("id", "s", "t", Role.TOOL, "c"))
    True
    """
    # Access __dict__ directly to bypass @deal.inv wrappers on property access.
    # Hypothesis may generate LogEntry objects with invalid state that would trigger
    # InvContractError on normal property access.
    d = log.__dict__
    RoleCls = Role  # Capture in local for faster access

    # Check required fields are non-empty
    session_id = d.get("session_id", "")
    timestamp = d.get("timestamp", "")
    if not (session_id and timestamp):
        return False

    # Check tool fields only valid for TOOL role
    role = d.get("role")
    if role != RoleCls.TOOL:
        if (
            d.get("tool_input") is not None
            or d.get("tool_status") is not None
            or d.get("tool_error") is not None
        ):
            return False

    # Check tool_status not empty string
    tool_status = d.get("tool_status")
    if tool_status is not None and tool_status == "":
        return False

    # Check tool_error not empty string
    tool_error = d.get("tool_error")
    if tool_error is not None and tool_error == "":
        return False

    return True


@deal.pre(lambda logs: len(logs) == 0 or all(_log_entry_is_valid(log) for log in logs))
@deal.post(
    lambda result: len(result) == 0 or all(len(s.session_id) > 0 for s in result)
)
@deal.post(lambda result: all(len(s.logs) > 0 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 "")
