"""
JARVIS — Database Layer
Async SQLAlchemy with SQLite.
Tables: sessions, conversations, embeddings_meta
"""

from datetime import datetime
from typing import Optional, List

from sqlalchemy import (
    Column, String, Text, DateTime, Integer,
    ForeignKey, Index, event
)
from sqlalchemy.ext.asyncio import (
    AsyncSession, create_async_engine, async_sessionmaker
)
from sqlalchemy.orm import DeclarativeBase, relationship
from sqlalchemy.future import select
from loguru import logger

from app.core.config import settings


# ─────────────────────────────────────────────
# Engine & session factory
# ─────────────────────────────────────────────
engine = create_async_engine(
    settings.DATABASE_URL,
    echo=settings.DEBUG,
    connect_args={"check_same_thread": False},
)

AsyncSessionLocal = async_sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,
)


# ─────────────────────────────────────────────
# Base model
# ─────────────────────────────────────────────
class Base(DeclarativeBase):
    pass


# ─────────────────────────────────────────────
# ORM Models
# ─────────────────────────────────────────────
class Session(Base):
    """A chat session (conversation thread)."""
    __tablename__ = "sessions"

    id         = Column(String(36), primary_key=True)   # UUID
    title      = Column(String(200), default="New Chat")
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    messages = relationship("Message", back_populates="session", cascade="all, delete-orphan")


class Message(Base):
    """A single message in a session."""
    __tablename__ = "messages"

    id         = Column(Integer, primary_key=True, autoincrement=True)
    session_id = Column(String(36), ForeignKey("sessions.id"), nullable=False)
    role       = Column(String(20), nullable=False)   # "user" | "assistant" | "system"
    content    = Column(Text, nullable=False)
    modality   = Column(String(20), default="text")   # text | image | voice | tool
    metadata_  = Column("metadata", Text, default="{}")  # JSON string
    created_at = Column(DateTime, default=datetime.utcnow)

    session = relationship("Session", back_populates="messages")

    __table_args__ = (
        Index("ix_messages_session_id", "session_id"),
        Index("ix_messages_created_at", "created_at"),
    )


class EmbeddingMeta(Base):
    """Metadata for stored FAISS embeddings."""
    __tablename__ = "embeddings_meta"

    id         = Column(Integer, primary_key=True, autoincrement=True)
    faiss_id   = Column(Integer, unique=True, nullable=False)   # index in FAISS
    session_id = Column(String(36), nullable=True)
    text       = Column(Text, nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow)


# ─────────────────────────────────────────────
# Init
# ─────────────────────────────────────────────
async def init_db():
    """Create all tables if they don't exist."""
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    logger.info("Database tables created/verified.")


# ─────────────────────────────────────────────
# Dependency
# ─────────────────────────────────────────────
async def get_db() -> AsyncSession:
    """FastAPI dependency — yields an async DB session."""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise


# ─────────────────────────────────────────────
# Helper CRUD functions
# ─────────────────────────────────────────────
async def create_session(db: AsyncSession, session_id: str, title: str = "New Chat") -> Session:
    sess = Session(id=session_id, title=title)
    db.add(sess)
    await db.flush()
    return sess


async def get_session(db: AsyncSession, session_id: str) -> Optional[Session]:
    result = await db.execute(select(Session).where(Session.id == session_id))
    return result.scalar_one_or_none()


async def list_sessions(db: AsyncSession, limit: int = 50) -> List[Session]:
    result = await db.execute(
        select(Session).order_by(Session.updated_at.desc()).limit(limit)
    )
    return result.scalars().all()


async def add_message(
    db: AsyncSession,
    session_id: str,
    role: str,
    content: str,
    modality: str = "text",
    metadata: str = "{}",
) -> Message:
    msg = Message(
        session_id=session_id,
        role=role,
        content=content,
        modality=modality,
        metadata_=metadata,
    )
    db.add(msg)
    await db.flush()
    # Update session timestamp
    sess = await get_session(db, session_id)
    if sess:
        sess.updated_at = datetime.utcnow()
    return msg


async def get_messages(
    db: AsyncSession,
    session_id: str,
    limit: int = 50,
) -> List[Message]:
    result = await db.execute(
        select(Message)
        .where(Message.session_id == session_id)
        .order_by(Message.created_at.asc())
        .limit(limit)
    )
    return result.scalars().all()


async def add_embedding_meta(
    db: AsyncSession,
    faiss_id: int,
    text: str,
    session_id: Optional[str] = None,
) -> EmbeddingMeta:
    em = EmbeddingMeta(faiss_id=faiss_id, text=text, session_id=session_id)
    db.add(em)
    await db.flush()
    return em