import logging
from typing import Optional, Set, Dict, Type
from weakref import WeakValueDictionary

logger = logging.getLogger(__name__)


class Session:
    """
    Session class to track model instances and manage database operations.

    Usage:
        with Session(db) as session:
            user = User(name="Alice")
            session.add(user)
            session.commit()

    Or:
        session = Session(db)
        user = User(name="Alice")
        session.add(user)
        session.commit()
        session.close()
    """

    _current_session: Optional["Session"] = None

    def __init__(self, db):
        """
        Initialize a new Session.

        Args:
            db: Database instance to use for this session
        """
        self.db = db
        self._new: Set = set()  # New objects to be inserted
        self._dirty: Set = set()  # Modified objects to be updated
        self._deleted: Set = set()  # Objects to be deleted
        self._identity_map: Dict[tuple, object] = (
            WeakValueDictionary()
        )  # (model_class, id) -> instance
        self._is_active = True  # Active by default for manual usage

    def __enter__(self):
        """Enter context manager - set this session as current."""
        self._is_active = True
        Session._current_session = self
        logger.info("Session started")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """Exit context manager - commit or rollback based on exception."""
        try:
            if exc_type:
                logger.error(
                    "Session exiting with exception, rolling back.",
                    exc_info=(exc_type, exc_value, traceback),
                )
                self.rollback()
                return False
            else:
                self.commit()
                logger.info("Session committed successfully")
        finally:
            self.close()

    def add(self, instance):
        """
        Add a model instance to the session for tracking.

        Args:
            instance: Model instance to track
        """
        if instance.id is None:
            # New object - needs INSERT
            self._new.add(instance)
            logger.debug(
                f"Added new instance to session: {instance.__class__.__name__}"
            )
        else:
            # Existing object - track for potential UPDATE
            self._mark_dirty(instance)
            # Add to identity map
            key = (instance.__class__, instance.id)
            self._identity_map[key] = instance

    def delete(self, instance):
        """
        Mark an instance for deletion.

        Args:
            instance: Model instance to delete
        """
        self._deleted.add(instance)
        # Remove from other sets if present
        self._new.discard(instance)
        self._dirty.discard(instance)
        logger.debug(
            f"Marked instance for deletion: {instance.__class__.__name__} id={instance.id}"
        )

    def _mark_dirty(self, instance):
        """Mark an instance as modified (needs UPDATE)."""
        if instance not in self._new:  # Don't mark new objects as dirty
            self._dirty.add(instance)

    def commit(self):
        """
        Commit all pending changes to the database.

        This will:
        1. INSERT all new objects
        2. UPDATE all dirty objects
        3. DELETE all deleted objects
        """
        if not self._is_active:
            raise RuntimeError("Cannot commit - session is not active")

        try:
            # Delete first
            for instance in self._deleted:
                self._execute_delete(instance)

            # Then insert new
            for instance in self._new:
                instance.save(self.db)
                # Add to identity map after save (now has ID)
                if instance.id is not None:
                    key = (instance.__class__, instance.id)
                    self._identity_map[key] = instance

            # Then update dirty
            for instance in self._dirty:
                if instance not in self._deleted:  # Skip if marked for deletion
                    instance.save(self.db)

            # Commit to database
            self.db.commit()

            # Clear tracking sets
            self._new.clear()
            self._dirty.clear()
            self._deleted.clear()

            logger.info("Session committed successfully")

        except Exception as e:
            logger.error(f"Error during commit: {e}", exc_info=True)
            self.rollback()
            raise

    def rollback(self):
        """
        Rollback all pending changes.

        Clears all tracked objects without saving to database.
        """
        self.db.backend.rollback()
        self._new.clear()
        self._dirty.clear()
        self._deleted.clear()
        logger.info("Session rolled back")

    def close(self):
        """Close the session and clear current session reference."""
        self._is_active = False
        if Session._current_session is self:
            Session._current_session = None
        self._new.clear()
        self._dirty.clear()
        self._deleted.clear()
        logger.info("Session closed")

    def get(self, model_class: Type, pk: int):
        """
        Get an instance by primary key, using identity map if available.

        Args:
            model_class: The model class to query
            pk: Primary key value

        Returns:
            Model instance or None
        """
        # Check identity map first
        key = (model_class, pk)
        if key in self._identity_map:
            logger.debug(f"Retrieved {model_class.__name__} id={pk} from identity map")
            return self._identity_map[key]

        # Load from database
        instance = model_class.get(pk)
        if instance:
            self._identity_map[key] = instance
            logger.debug(f"Loaded {model_class.__name__} id={pk} from database")
        return instance

    def query(self, model_class):
        """
        Create a new QuerySet for the given model class.

        Args:
            model_class: The model class to query

        Returns:
            QuerySet instance
        """
        from .query import QuerySet

        return QuerySet(model_class, self.db)

    @classmethod
    def get_current(cls) -> Optional["Session"]:
        """Get the current active session, if any."""
        return cls._current_session

    def _execute_delete(self, instance):
        """Execute DELETE SQL for an instance."""
        if instance.id is None:
            logger.warning(f"Cannot delete {instance.__class__.__name__} - no id set")
            return

        table_name = f'"{instance._table_name}"'
        sql = f"DELETE FROM {table_name} WHERE id = {self.db.backend.placeholder_char}"
        self.db.execute(sql, [instance.id])

        # Remove from session's identity map
        key = (instance.__class__, instance.id)
        self._identity_map.pop(key, None)

        # Also remove from database's global identity map
        self.db.remove_from_identity_map(instance)

        logger.info(f"Deleted {instance._table_name} with id={instance.id}")

    def flush(self):
        """
        Flush pending changes to database without committing.

        This executes all SQL but doesn't call commit on the backend,
        allowing for transaction-level control.
        """
        # Execute deletes
        for instance in self._deleted:
            self._execute_delete(instance)

        # Execute inserts
        for instance in self._new:
            instance.save(self.db)
            if instance.id is not None:
                key = (instance.__class__, instance.id)
                self._identity_map[key] = instance

        # Execute updates
        for instance in self._dirty:
            if instance not in self._deleted:
                instance.save(self.db)

        logger.info("Session flushed (changes executed but not committed)")
