"""SAM primary session -- multi-session management on single I2P tunnel.

Ported from net.i2p.sam.SAMv3Handler PRIMARY session support.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from i2p_sam.sessions_db import SessionRecord

logger = logging.getLogger(__name__)


class PrimarySession:
    """Manages multiple subsessions (STREAM, DATAGRAM, RAW) on one I2P session.

    A PRIMARY session allows a single I2P destination to be used for
    multiple communication styles simultaneously. Subsessions are keyed
    by their FROM_PORT to allow multiplexing.
    """

    def __init__(self, nickname: str, destination_b64: str) -> None:
        self._nickname = nickname
        self._destination_b64 = destination_b64
        self._subsessions: dict[str, "SessionRecord"] = {}

    async def add_subsession(self, sub_key: str, style: str,
                             handler: object) -> "SessionRecord":
        """Add a subsession under this primary session.

        Args:
            sub_key: Subsession identifier (typically FROM_PORT value).
            style: Communication style (STREAM, DATAGRAM, RAW).
            handler: The SAMHandler managing this subsession.

        Returns:
            The created SessionRecord.

        Raises:
            ValueError: If sub_key already exists.
        """
        from i2p_sam.sessions_db import SessionRecord
        from i2p_sam.utils import generate_transient_destination

        if sub_key in self._subsessions:
            raise ValueError(f"Subsession {sub_key} already exists")

        raw, b64 = generate_transient_destination()
        record = SessionRecord(
            nickname=f"{self._nickname}:{sub_key}",
            style=style,
            destination=raw,
            destination_b64=b64,
            handler=handler,  # type: ignore[arg-type]
        )
        self._subsessions[sub_key] = record
        logger.info("Added subsession %s (style=%s) to primary %s",
                     sub_key, style, self._nickname)
        return record

    async def remove_subsession(self, sub_key: str) -> bool:
        """Remove a subsession.

        Args:
            sub_key: Subsession identifier to remove.

        Returns:
            True if removed, False if not found.
        """
        if sub_key in self._subsessions:
            del self._subsessions[sub_key]
            logger.info("Removed subsession %s from primary %s",
                         sub_key, self._nickname)
            return True
        return False

    def get_subsession(self, sub_key: str) -> "SessionRecord | None":
        """Get a subsession by key.

        Args:
            sub_key: Subsession identifier.

        Returns:
            The SessionRecord, or None if not found.
        """
        return self._subsessions.get(sub_key)

    @property
    def subsession_count(self) -> int:
        """Number of active subsessions."""
        return len(self._subsessions)