from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, scoped_session, relationship
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, create_engine
from datetime import datetime
from typing import Optional, Dict, List, Any
import os


Base = declarative_base()


class ORMManager:
    """Enhanced ORM manager with SQLAlchemy support"""
    
    def __init__(self, db_url: Optional[str] = None):
        if db_url is None:
            db_url = os.environ.get('DATABASE_URL', 'sqlite:///app.db')
        
        self.engine = create_engine(db_url)
        self.Session = scoped_session(sessionmaker(bind=self.engine))
        self.base = Base
    
    def create_all(self):
        """Create all tables"""
        Base.metadata.create_all(self.engine)
    
    def drop_all(self):
        """Drop all tables"""
        Base.metadata.drop_all(self.engine)
    
    def get_session(self):
        """Get a database session"""
        return self.Session()
    
    def close_session(self):
        """Close the current session"""
        self.Session.remove()


class QueryBuilder:
    """Type-safe query builder"""
    
    def __init__(self, model_class, session):
        self.model_class = model_class
        self.session = session
        self._query = session.query(model_class)
    
    def filter(self, **kwargs):
        """Filter results"""
        for key, value in kwargs.items():
            column = getattr(self.model_class, key)
            self._query = self._query.filter(column == value)
        return self
    
    def filter_by(self, **kwargs):
        """Filter by keyword arguments"""
        self._query = self._query.filter_by(**kwargs)
        return self
    
    def order_by(self, column_name: str, desc: bool = False):
        """Order results"""
        column = getattr(self.model_class, column_name)
        if desc:
            self._query = self._query.order_by(column.desc())
        else:
            self._query = self._query.order_by(column)
        return self
    
    def limit(self, count: int):
        """Limit results"""
        self._query = self._query.limit(count)
        return self
    
    def offset(self, count: int):
        """Offset results"""
        self._query = self._query.offset(count)
        return self
    
    def all(self):
        """Get all results"""
        return self._query.all()
    
    def first(self):
        """Get first result"""
        return self._query.first()
    
    def count(self):
        """Count results"""
        return self._query.count()
    
    def one(self):
        """Get exactly one result"""
        return self._query.one()
    
    def one_or_none(self):
        """Get one result or None"""
        return self._query.one_or_none()


class BaseModel(Base):
    """Base model class with common fields"""
    __abstract__ = True
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    @classmethod
    def create(cls, session, **kwargs):
        """Create a new instance"""
        instance = cls(**kwargs)
        session.add(instance)
        session.commit()
        return instance
    
    @classmethod
    def get_by_id(cls, session, record_id: int):
        """Get by ID"""
        return session.query(cls).filter_by(id=record_id).first()
    
    @classmethod
    def query(cls, session):
        """Get a query builder"""
        return QueryBuilder(cls, session)
    
    def update(self, session, **kwargs):
        """Update instance"""
        for key, value in kwargs.items():
            setattr(self, key, value)
        session.commit()
    
    def delete(self, session):
        """Delete instance"""
        session.delete(self)
        session.commit()
    
    def to_dict(self):
        """Convert to dictionary"""
        return {
            column.name: getattr(self, column.name)
            for column in self.__table__.columns
        }


orm = ORMManager()
