from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime, ForeignKey, Boolean, Enum as SQLEnum
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship, Session
from datetime import datetime
from typing import Optional
import enum
import os
from dotenv import load_dotenv

load_dotenv()

Base = declarative_base()

# Database connection
DATABASE_URL = f"postgresql+psycopg2://{os.getenv('pgql_user')}:{os.getenv('pgql_psw')}@localhost:5432/{os.getenv('pgql_db')}"
engine = create_engine(DATABASE_URL, pool_size=10, max_overflow=20)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


class PermissionLevel(enum.Enum):
    READ = "READ"
    WRITE = "WRITE"
    SHARE = "SHARE"  # Write + Share permissions
    ADMIN = "ADMIN"


class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True, index=True)
    email = Column(String(255), unique=True, index=True, nullable=False)
    username = Column(String(100), unique=True, index=True, nullable=False)
    password_hash = Column(String(255), nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow)

    # Relationships
    documents = relationship("Document", back_populates="owner", foreign_keys="Document.user_id")
    notebooks = relationship("Notebook", back_populates="owner")
    chat_sessions = relationship("ChatSession", back_populates="user")
    shared_with_me = relationship("DocumentShare", foreign_keys="DocumentShare.shared_with_user_id", back_populates="shared_with_user")
    shared_by_me = relationship("DocumentShare", foreign_keys="DocumentShare.owner_id", back_populates="owner")


class Document(Base):
    __tablename__ = "documents"

    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
    filename = Column(String(500), nullable=False)
    original_filename = Column(String(500), nullable=False)
    llamacloud_file_id = Column(String(255), nullable=True)
    pipeline_id = Column(String(255), nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)

    # Relationships
    owner = relationship("User", back_populates="documents", foreign_keys=[user_id])
    document_summaries = relationship("DocumentSummary", back_populates="document")
    notebook_links = relationship("NotebookDocument", back_populates="document")
    shares = relationship("DocumentShare", back_populates="document")


class Notebook(Base):
    """Collection of documents (like Google NotebookLM)"""
    __tablename__ = "notebooks"

    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
    name = Column(String(500), nullable=False)
    description = Column(Text, nullable=True)
    podcast_path = Column(String(500), nullable=True)
    podcast_generated_at = Column(DateTime, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Relationships
    owner = relationship("User", back_populates="notebooks")
    document_links = relationship("NotebookDocument", back_populates="notebook", cascade="all, delete-orphan")
    chat_sessions = relationship("ChatSession", back_populates="notebook")


class NotebookDocument(Base):
    """Junction table linking notebooks to documents"""
    __tablename__ = "notebook_documents"

    id = Column(Integer, primary_key=True, index=True)
    notebook_id = Column(Integer, ForeignKey("notebooks.id"), nullable=False, index=True)
    document_id = Column(Integer, ForeignKey("documents.id"), nullable=False, index=True)
    display_order = Column(Integer, default=0)
    added_at = Column(DateTime, default=datetime.utcnow)

    # Relationships
    notebook = relationship("Notebook", back_populates="document_links")
    document = relationship("Document", back_populates="notebook_links")


class DocumentSummary(Base):
    """Summary/analysis of individual document (renamed from old Notebook)"""
    __tablename__ = "document_summaries"

    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
    document_id = Column(Integer, ForeignKey("documents.id"), nullable=False, index=True)
    summary = Column(Text, nullable=True)
    highlights = Column(Text, nullable=True)  # JSON array stored as text
    questions = Column(Text, nullable=True)  # JSON array stored as text
    answers = Column(Text, nullable=True)  # JSON array stored as text
    md_content = Column(Text, nullable=True)
    mind_map_path = Column(String(500), nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)

    # Relationships
    document = relationship("Document", back_populates="document_summaries")


class ChatSession(Base):
    __tablename__ = "chat_sessions"

    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
    document_id = Column(Integer, ForeignKey("documents.id"), nullable=True, index=True)  # Made nullable
    notebook_id = Column(Integer, ForeignKey("notebooks.id"), nullable=True, index=True)  # NEW: for notebook-level chats
    title = Column(String(500), nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Relationships
    user = relationship("User", back_populates="chat_sessions")
    notebook = relationship("Notebook", back_populates="chat_sessions")
    messages = relationship("ChatMessage", back_populates="session", cascade="all, delete-orphan")


class ChatMessage(Base):
    __tablename__ = "chat_messages"

    id = Column(Integer, primary_key=True, index=True)
    session_id = Column(Integer, ForeignKey("chat_sessions.id"), nullable=False, index=True)
    role = Column(String(50), nullable=False)  # 'user' or 'assistant'
    content = Column(Text, nullable=False)
    sources = Column(Text, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)

    # Relationships
    session = relationship("ChatSession", back_populates="messages")


class DocumentShare(Base):
    __tablename__ = "document_shares"

    id = Column(Integer, primary_key=True, index=True)
    document_id = Column(Integer, ForeignKey("documents.id"), nullable=True, index=True)
    notebook_id = Column(Integer, ForeignKey("notebooks.id"), nullable=True, index=True)  # NEW: share notebooks too
    owner_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
    shared_with_user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
    permission_level = Column(SQLEnum(PermissionLevel), nullable=False, default=PermissionLevel.READ)
    created_at = Column(DateTime, default=datetime.utcnow)

    # Relationships
    document = relationship("Document", back_populates="shares")
    owner = relationship("User", foreign_keys=[owner_id], back_populates="shared_by_me")
    shared_with_user = relationship("User", foreign_keys=[shared_with_user_id], back_populates="shared_with_me")


# Helper functions
def get_db() -> Session:
    """Get database session"""
    db = SessionLocal()
    try:
        return db
    finally:
        pass


def init_db():
    """Initialize database tables"""
    Base.metadata.create_all(bind=engine)


def drop_all_tables():
    """Drop all tables (use with caution!)"""
    Base.metadata.drop_all(bind=engine)