data_store.py

import os
import json
import datetime
from typing import Dict, Any, List, Optional
from sqlalchemy import (
    create_engine, Column, Integer, String, Text,
    DateTime, Boolean, ForeignKey, BigInteger, LargeBinary, event,
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship
from sqlalchemy.exc import SQLAlchemyError
from dotenv import load_dotenv

load_dotenv()

ENV = os.getenv("FLASK_ENV", "development")
if ENV == "production":
    DATABASE_URL = os.getenv("DATABASE_URL", "mysql://user:password@localhost/drive_db")
else:
    DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///drive.db")

if DATABASE_URL.startswith("sqlite"):
    engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False}, echo=False)
else:
    engine = create_engine(DATABASE_URL, echo=False, pool_size=5, max_overflow=10, pool_timeout=30)

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()








class User(Base):
    __tablename__ = "users"

    id           = Column(Integer, primary_key=True)
    username     = Column(String(255), unique=True, nullable=False)
    email        = Column(String(255), unique=True, nullable=False)
    password     = Column(String(255), nullable=False)
    role         = Column(String(50), default="subuser")   # "admin" | "subuser"
    is_active    = Column(Boolean, default=True)
    storage_quota= Column(BigInteger, default=5 * 1024 * 1024 * 1024)  # 5 GB default
    storage_used = Column(BigInteger, default=0)
    created_at   = Column(DateTime, default=datetime.datetime.utcnow)
    updated_at   = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow)

    folders = relationship("Folder", back_populates="owner", cascade="all, delete-orphan")
    files   = relationship("DriveFile", back_populates="owner", cascade="all, delete-orphan")

    def to_dict(self):
        return {
            "id": self.id,
            "username": self.username,
            "email": self.email,
            "role": self.role,
            "is_active": self.is_active,
            "storage_quota": self.storage_quota,
            "storage_used": self.storage_used,
            "created_at": self.created_at.isoformat() if self.created_at else None,
        }


class Folder(Base):
    __tablename__ = "folders"

    id         = Column(Integer, primary_key=True)
    name       = Column(String(255), nullable=False)
    parent_id  = Column(Integer, ForeignKey("folders.id"), nullable=True)
    owner_id   = Column(Integer, ForeignKey("users.id"), nullable=False)
    created_at = Column(DateTime, default=datetime.datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow)

    owner    = relationship("User", back_populates="folders")
    children = relationship("Folder", backref="parent", remote_side=[id], foreign_keys=[parent_id])
    files    = relationship("DriveFile", back_populates="folder", cascade="all, delete-orphan")

    def to_dict(self):
        return {
            "id": self.id,
            "name": self.name,
            "parent_id": self.parent_id,
            "owner_id": self.owner_id,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
        }


class DriveFile(Base):
    __tablename__ = "drive_files"

    id                = Column(Integer, primary_key=True)
    filename          = Column(String(255), nullable=False)       # UUID-based disk name
    original_filename = Column(String(255), nullable=False)
    mime_type         = Column(String(200))
    file_size         = Column(BigInteger, default=0)
    folder_id         = Column(Integer, ForeignKey("folders.id"), nullable=True)
    owner_id          = Column(Integer, ForeignKey("users.id"), nullable=False)
    is_text           = Column(Boolean, default=False)             # created via text editor
    text_content      = Column(Text, nullable=True)                # inline text content
    file_data         = Column(LargeBinary, nullable=True)         # binary file content (DB storage)
    share_token       = Column(String(64), nullable=True, unique=True)
    is_shared         = Column(Boolean, default=False)
    share_expires_at  = Column(DateTime, nullable=True)
    created_at        = Column(DateTime, default=datetime.datetime.utcnow)
    updated_at        = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow)

    owner  = relationship("User", back_populates="files")
    folder = relationship("Folder", back_populates="files")

    def to_dict(self):
        return {
            "id": self.id,
            "filename": self.filename,
            "original_filename": self.original_filename,
            "mime_type": self.mime_type,
            "file_size": self.file_size,
            "folder_id": self.folder_id,
            "owner_id": self.owner_id,
            "is_text": self.is_text,
            "text_content": self.text_content,
            "share_token": self.share_token,
            "is_shared": self.is_shared,
            "share_expires_at": self.share_expires_at.isoformat() if self.share_expires_at else None,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
        }



class DriveDBManager:
    def __init__(self):
        Base.metadata.create_all(bind=engine)

    def get_session(self):
        return SessionLocal()

    def is_setup_complete(self) -> bool:
        db = self.get_session()
        try:
            return db.query(User).filter(User.role == "admin").count() > 0
        except SQLAlchemyError:
            return False
        finally:
            db.close()
            
            
            
            
            
            
            
            
            



    # ── Users ──────────────────────────────────

    def create_user(self, data: dict) -> Optional[User]:
        db = self.get_session()
        try:
            user = User(**data)
            db.add(user)
            db.commit()
            db.refresh(user)
            return user
        except SQLAlchemyError as e:
            db.rollback()
            print(f"create_user error: {e}")
            return None
        finally:
            db.close()

    def get_user_by_username(self, username: str) -> Optional[dict]:
        db = self.get_session()
        try:
            u = db.query(User).filter(User.username == username).first()
            return u.to_dict() | {"password": u.password} if u else None
        finally:
            db.close()

    def get_user_by_id(self, user_id: int) -> Optional[dict]:
        db = self.get_session()
        try:
            u = db.query(User).filter(User.id == user_id).first()
            return u.to_dict() if u else None
        finally:
            db.close()

    def get_all_users(self) -> List[dict]:
        db = self.get_session()
        try:
            return [u.to_dict() for u in db.query(User).filter(User.role != "admin").all()]
        finally:
            db.close()

    _UPDATABLE_USER_FIELDS = {"email", "password", "is_active", "storage_quota", "username"}

    def update_user(self, user_id: int, data: dict) -> bool:
        db = self.get_session()
        try:
            safe = {k: v for k, v in data.items() if k in self._UPDATABLE_USER_FIELDS}
            db.query(User).filter(User.id == user_id).update(safe)
            db.commit()
            return True
        except SQLAlchemyError as e:
            db.rollback()
            print(f"update_user error: {e}")
            return False
        finally:
            db.close()

    def delete_user(self, user_id: int) -> bool:
        db = self.get_session()
        try:
            u = db.query(User).filter(User.id == user_id).first()
            if u:
                db.delete(u)
                db.commit()
                return True
            return False
        except SQLAlchemyError as e:
            db.rollback()
            return False
        finally:
            db.close()

    def update_storage_used(self, user_id: int, delta: int):
        """Add or subtract bytes from user's used storage."""
        db = self.get_session()
        try:
            u = db.query(User).filter(User.id == user_id).first()
            if u:
                u.storage_used = max(0, (u.storage_used or 0) + delta)
                db.commit()
        except SQLAlchemyError:
            db.rollback()
        finally:
            db.close()

    # ── Folders ────────────────────────────────

    def create_folder(self, name: str, owner_id: int, parent_id: Optional[int] = None) -> Optional[dict]:
        db = self.get_session()
        try:
            f = Folder(name=name, owner_id=owner_id, parent_id=parent_id)
            db.add(f)
            db.commit()
            db.refresh(f)
            return f.to_dict()
        except SQLAlchemyError as e:
            db.rollback()
            print(f"create_folder error: {e}")
            return None
        finally:
            db.close()

    def get_folders(self, owner_id: int, parent_id: Optional[int] = None, is_admin: bool = False) -> List[dict]:
        db = self.get_session()
        try:
            q = db.query(Folder)
            if not is_admin:
                q = q.filter(Folder.owner_id == owner_id)
            q = q.filter(Folder.parent_id == parent_id)
            return [f.to_dict() for f in q.all()]
        finally:
            db.close()

    def get_folder(self, folder_id: int) -> Optional[dict]:
        db = self.get_session()
        try:
            f = db.query(Folder).filter(Folder.id == folder_id).first()
            return f.to_dict() if f else None
        finally:
            db.close()

    def rename_folder(self, folder_id: int, name: str, requestor_id: int, is_admin: bool) -> bool:
        db = self.get_session()
        try:
            f = db.query(Folder).filter(Folder.id == folder_id).first()
            if not f:
                return False
            if not is_admin and f.owner_id != requestor_id:
                return False
            f.name = name
            db.commit()
            return True
        except SQLAlchemyError:
            db.rollback()
            return False
        finally:
            db.close()

    def delete_folder(self, folder_id: int, requestor_id: int, is_admin: bool) -> bool:
        db = self.get_session()
        try:
            f = db.query(Folder).filter(Folder.id == folder_id).first()
            if not f:
                return False
            if not is_admin and f.owner_id != requestor_id:
                return False
            # collect all disk files before deleting
            disk_files = self._collect_files_in_folder(db, folder_id)
            db.delete(f)
            db.commit()
            return disk_files  # caller handles disk deletion
        except SQLAlchemyError:
            db.rollback()
            return False
        finally:
            db.close()

    def _collect_files_in_folder(self, db, folder_id: int) -> List[str]:
        result = []
        files = db.query(DriveFile).filter(DriveFile.folder_id == folder_id).all()
        for f in files:
            result.append(f.filename)
        sub_folders = db.query(Folder).filter(Folder.parent_id == folder_id).all()
        for sf in sub_folders:
            result.extend(self._collect_files_in_folder(db, sf.id))
        return result

    def get_folder_breadcrumb(self, folder_id: Optional[int]) -> List[dict]:
        """Return list of ancestors from root → current."""
        if not folder_id:
            return []
        db = self.get_session()
        crumbs = []
        try:
            fid = folder_id
            while fid:
                f = db.query(Folder).filter(Folder.id == fid).first()
                if not f:
                    break
                crumbs.insert(0, {"id": f.id, "name": f.name})
                fid = f.parent_id
        finally:
            db.close()
        return crumbs

    # ── Files ──────────────────────────────────

    def add_file(self, data: dict) -> Optional[dict]:
        db = self.get_session()
        try:
            f = DriveFile(**data)
            db.add(f)
            db.commit()
            db.refresh(f)
            return f.to_dict()
        except SQLAlchemyError as e:
            db.rollback()
            print(f"add_file error: {e}")
            return None
        finally:
            db.close()

    def get_files(self, owner_id: int, folder_id: Optional[int] = None, is_admin: bool = False) -> List[dict]:
        db = self.get_session()
        try:
            q = db.query(DriveFile)
            if not is_admin:
                q = q.filter(DriveFile.owner_id == owner_id)
            q = q.filter(DriveFile.folder_id == folder_id)
            return [f.to_dict() for f in q.order_by(DriveFile.created_at.desc()).all()]
        finally:
            db.close()

    def get_file(self, file_id: int) -> Optional[dict]:
        db = self.get_session()
        try:
            f = db.query(DriveFile).filter(DriveFile.id == file_id).first()
            return f.to_dict() if f else None
        finally:
            db.close()

    def get_file_data(self, file_id: int) -> Optional[bytes]:
        """Fetch only the binary content of a file — avoids loading blob in list queries."""
        db = self.get_session()
        try:
            f = db.query(DriveFile).filter(DriveFile.id == file_id).first()
            return f.file_data if f else None
        finally:
            db.close()

    def get_file_by_token(self, token: str) -> Optional[dict]:
        db = self.get_session()
        try:
            f = db.query(DriveFile).filter(DriveFile.share_token == token, DriveFile.is_shared == True).first()
            if not f:
                return None
            fd = f.to_dict()
            # check expiry
            if f.share_expires_at and f.share_expires_at < datetime.datetime.utcnow():
                return None
            return fd
        finally:
            db.close()

    _UPDATABLE_FILE_FIELDS = {"text_content", "original_filename", "mime_type", "folder_id"}

    def update_file(self, file_id: int, data: dict, requestor_id: int, is_admin: bool) -> bool:
        db = self.get_session()
        try:
            f = db.query(DriveFile).filter(DriveFile.id == file_id).first()
            if not f:
                return False
            if not is_admin and f.owner_id != requestor_id:
                return False
            for k, v in data.items():
                if k not in self._UPDATABLE_FILE_FIELDS:
                    raise ValueError(f"Field '{k}' is not updatable")
                setattr(f, k, v)
            db.commit()
            return True
        except SQLAlchemyError:
            db.rollback()
            return False
        finally:
            db.close()

    def delete_file(self, file_id: int, requestor_id: int, is_admin: bool):
        """Returns the on-disk filename and owner_id on success, or None."""
        db = self.get_session()
        try:
            f = db.query(DriveFile).filter(DriveFile.id == file_id).first()
            if not f:
                return None
            if not is_admin and f.owner_id != requestor_id:
                return None
            info = {"filename": f.filename, "owner_id": f.owner_id, "file_size": f.file_size}
            db.delete(f)
            db.commit()
            return info
        except SQLAlchemyError:
            db.rollback()
            return None
        finally:
            db.close()

    def search_files(self, owner_id: int, query: str, is_admin: bool) -> List[dict]:
        db = self.get_session()
        try:
            q = db.query(DriveFile)
            if not is_admin:
                q = q.filter(DriveFile.owner_id == owner_id)
            q = q.filter(DriveFile.original_filename.ilike(f"%{query}%"))
            return [f.to_dict() for f in q.limit(50).all()]
        finally:
            db.close()

    def set_share(self, file_id: int, token: str, is_shared: bool, expires_at=None, requestor_id: int = 0, is_admin: bool = False) -> bool:
        db = self.get_session()
        try:
            f = db.query(DriveFile).filter(DriveFile.id == file_id).first()
            if not f:
                return False
            if not is_admin and f.owner_id != requestor_id:
                return False
            f.share_token = token
            f.is_shared = is_shared
            f.share_expires_at = expires_at
            db.commit()
            return True
        except SQLAlchemyError:
            db.rollback()
            return False
        finally:
            db.close()

    def get_all_files_for_admin(self) -> List[dict]:
        db = self.get_session()
        try:
            return [f.to_dict() for f in db.query(DriveFile).order_by(DriveFile.created_at.desc()).all()]
        finally:
            db.close()


db_manager = DriveDBManager()