class AgentConnection extends Data.TaggedClass("AgentConnection")<
  ConnectionBase & { readonly auth: AgentContext }
> {
  private readonly brandValue!: never;
}

export class AgentContext extends Data.TaggedClass("AgentContext")<{
  readonly agentId: AgentId;
  readonly agentStatus: AgentStatus;
  readonly ownerUserId: UserId;
}> {}

export function agentContextFrom(parts: {
  readonly agentId: AgentId;
  readonly agentStatus: string;
  readonly ownerUserId: UserId;
}): Effect.Effect<AgentContext>

export type AgentStatus = "active" | "suspended";

export type Connection = UnauthenticatedConnection | AgentConnection;

export class ConnectionManager {
  /**
   * Connections and their per-agent delivery projection share one Ref so
   * authentication, disconnect cleanup, and subscription updates are atomic.
   * This prevents an old last-disconnect cleanup from deleting a newly
   * authenticated socket's freshly hydrated subscriptions.
   */
  private readonly stateRef: Ref.Ref<ConnectionManagerState> = Effect.runSync(
    Ref.make({
      connections: HashMap.empty<ConnectionId, Connection>(),
      agentConversationSubscriptions: HashMap.empty<
        AgentId,
        HashSet.HashSet<ConversationId>
      >(),
    }),
  );

  /**
   * Insert a fresh `UnauthenticatedConnection`. Called by the socket handler
   * at WebSocket open. The Connect handler promotes it to the agent arm.
   * @param connId Value supplied to the operation.
   * @param socket Value supplied to the operation.
   * @param originator Value supplied to the operation.
   * @returns The add unauthenticated result.
   */
  addUnauthenticated(
    connId: ConnectionId,
    socket: WebSocketRef,
    originator: Originator,
  ): Effect.Effect<void> {
    return Ref.update(this.stateRef, (state) => ({
      ...state,
      connections: HashMap.set(
        state.connections,
        connId,
        new UnauthenticatedConnection({ connId, socket, originator }),
      ),
    }));
  }

  /**
   * Non-mutating read. Callers discriminate on the returned arm's `_tag`.
   * @param connId Value supplied to the operation.
   * @returns The current result.
   */
  peek(connId: ConnectionId): Effect.Effect<Option.Option<Connection>> {
    return Ref.get(this.stateRef).pipe(
      Effect.map((state) => HashMap.get(state.connections, connId)),
    );
  }

  /**
   * Snapshot of every connection arm. Callers iterate + discriminate on `_tag`
   * (e.g. The shutdown loop reads `arm.socket.shutdown`).
   * @returns The current result.
   */
  allConnections(): Effect.Effect<readonly Connection[]> {
    return Ref.get(this.stateRef).pipe(
      Effect.map((state) => Array.from(HashMap.values(state.connections))),
    );
  }

  /**
   * Current connection count.
   * @returns The current result.
   */
  currentSize(): Effect.Effect<number> {
    return Ref.get(this.stateRef).pipe(
      Effect.map((state) => HashMap.size(state.connections)),
    );
  }

  /**
   * Atomic per-connection authentication gate. Mints the agent arm from the
   * unauthenticated entry and returns a `TransitionOutcome` whose success arm
   * carries the minted connection, so callers narrow without a cast.
   * @param connId Value supplied to the operation.
   * @param auth Value supplied to the operation.
   * @returns The current result.
   */
  authenticate(
    connId: ConnectionId,
    auth: AgentContext,
  ): Effect.Effect<TransitionOutcome> {
    return Ref.modify(this.stateRef, (state) => {
      const current = HashMap.get(state.connections, connId);
      if (Option.isNone(current)) {
        return [{ kind: "not-connected" } as const, state];
      }
      return Match.value(current.value).pipe(
        Match.tag(
          "AgentConnection",
          (existing): [TransitionOutcome, typeof state] => [
            { kind: "already-connected", existing },
            state,
          ],
        ),
        Match.tag(
          "UnauthenticatedConnection",
          (unauth): [TransitionOutcome, typeof state] => {
            const authed = new AgentConnection({
              connId: unauth.connId,
              socket: unauth.socket,
              originator: unauth.originator,
              auth,
            });
            return [
              { kind: "ok-agent", authed },
              {
                ...state,
                connections: HashMap.set(state.connections, connId, authed),
              },
            ];
          },
        ),
        Match.exhaustive,
      );
    });
  }

export const connectionManagerLive = Layer.sync(
  ConnectionManagerTag,
  () => new ConnectionManager(),
)

export class ConnectionManagerTag extends Context.Tag(
  "moltzap/ConnectionManager",
)<ConnectionManagerTag, ConnectionManager>() {}

export class ConnectionTag extends Context.Tag("moltzap/Connection")<
  ConnectionTag,
  Connection
>() {}

export type Originator = ReverseClient;

export type PrincipalBoundaryCanaries = [
  UnauthenticatedHasNoAuth,
  ForgedAgentRejected,
  InvalidBootPhaseRejected,
];

export const principalCanaryRefs: readonly unknown[] = [
  agentIdValue,
  principalTag,
  narrowOutcome,
  bootFail,
] as const

export type TransitionOutcome =
  | { readonly kind: "not-connected" }

class UnauthenticatedConnection extends Data.TaggedClass(
  "UnauthenticatedConnection",
)<ConnectionBase> {
  private readonly brand!: never;
}

export interface WebSocketRef {
  /**
   * Write a raw frame to this connection. Fails with SocketError on send
   * failure or if the socket is already closed.
   */
  readonly write: (raw: string) => Effect.Effect<void, SocketError>;
  /** Close this connection's scope, tearing down the underlying socket. */
  readonly shutdown: Effect.Effect<void>;
}