from datetime import datetime
from typing import Any

from pgvector.sqlalchemy import Vector
from pydantic import BaseModel
from sqlalchemy import (
    Boolean,
    DateTime,
    ForeignKey,
    Index,
    Integer,
    Text,
    UniqueConstraint,
    func,
    text,
)
from sqlalchemy.dialects.postgresql import JSONB, TSVECTOR
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column


class Base(DeclarativeBase):
    pass


class SessionModel(Base):
    __tablename__ = "sessions"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    host: Mapped[str] = mapped_column(Text, nullable=False)
    tool: Mapped[str] = mapped_column(Text, nullable=False)
    session_id: Mapped[str] = mapped_column(Text, nullable=False)
    title: Mapped[str] = mapped_column(Text, nullable=False, server_default=text("'Untitled'"))
    working_directory: Mapped[str] = mapped_column(Text, nullable=False, server_default=text("''"))
    model_name: Mapped[str | None] = mapped_column(Text, nullable=True)
    created_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    last_activity_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    message_count: Mapped[int] = mapped_column(Integer, nullable=False, server_default=text("0"))
    hidden: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default=text("false"))
    metadata_: Mapped[dict[str, Any]] = mapped_column(
        "metadata", JSONB, nullable=False, server_default=text("'{}'::jsonb")
    )
    synced_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    content_tsv: Mapped[str | None] = mapped_column(TSVECTOR, nullable=True)
    embedding: Mapped[list[float] | None] = mapped_column(Vector(1536), nullable=True)
    embedding_model: Mapped[str | None] = mapped_column(Text, nullable=True)
    embedded_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)

    __table_args__ = (
        UniqueConstraint("host", "tool", "session_id", name="uq_sessions_host_tool_session_id"),
        Index("ix_sessions_tool", "tool"),
        Index("ix_sessions_host", "host"),
        Index("ix_sessions_last_activity_at", text("last_activity_at DESC")),
        Index("ix_sessions_working_directory", "working_directory"),
        Index("idx_sessions_content_tsv", "content_tsv", postgresql_using="gin"),
        Index(
            "idx_sessions_embedding",
            "embedding",
            postgresql_using="hnsw",
            postgresql_ops={"embedding": "vector_cosine_ops"},
            postgresql_with={"m": 16, "ef_construction": 64},
        ),
    )


class MessageModel(Base):
    __tablename__ = "messages"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    session_pk: Mapped[int] = mapped_column(
        Integer, ForeignKey("sessions.id", ondelete="CASCADE"), nullable=False
    )
    host: Mapped[str] = mapped_column(Text, nullable=False)
    tool: Mapped[str] = mapped_column(Text, nullable=False)
    session_id: Mapped[str] = mapped_column(Text, nullable=False)
    node_id: Mapped[str] = mapped_column(Text, nullable=False)
    parent_node_id: Mapped[str | None] = mapped_column(Text, nullable=True)
    message_id: Mapped[str] = mapped_column(Text, nullable=False)
    role: Mapped[str] = mapped_column(Text, nullable=False)
    content: Mapped[str] = mapped_column(Text, nullable=False, server_default=text("''"))
    created_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    metadata_: Mapped[dict[str, Any]] = mapped_column(
        "metadata", JSONB, nullable=False, server_default=text("'{}'::jsonb")
    )
    synced_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    content_tsv: Mapped[str | None] = mapped_column(TSVECTOR, nullable=True)
    embedding: Mapped[list[float] | None] = mapped_column(Vector(1536), nullable=True)
    embedding_model: Mapped[str | None] = mapped_column(Text, nullable=True)
    embedded_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)

    __table_args__ = (
        UniqueConstraint(
            "host", "tool", "session_id", "node_id", name="uq_messages_host_tool_session_node"
        ),
        Index("ix_messages_session_pk", "session_pk"),
        Index("ix_messages_role", "role"),
        Index("ix_messages_created_at", "created_at"),
        Index("ix_messages_host_tool_session_id", "host", "tool", "session_id"),
        Index("idx_messages_content_tsv", "content_tsv", postgresql_using="gin"),
        Index(
            "idx_messages_content_trgm",
            "content",
            postgresql_using="gin",
            postgresql_ops={"content": "gin_trgm_ops"},
        ),
        Index(
            "idx_messages_embedding",
            "embedding",
            postgresql_using="hnsw",
            postgresql_ops={"embedding": "vector_cosine_ops"},
            postgresql_with={"m": 16, "ef_construction": 64},
        ),
    )


class SyncLogModel(Base):
    __tablename__ = "sync_log"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    host: Mapped[str] = mapped_column(Text, nullable=False)
    tool: Mapped[str | None] = mapped_column(Text, nullable=True)
    started_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    sessions_upserted: Mapped[int] = mapped_column(
        Integer, nullable=False, server_default=text("0")
    )
    messages_upserted: Mapped[int] = mapped_column(
        Integer, nullable=False, server_default=text("0")
    )
    status: Mapped[str] = mapped_column(Text, nullable=False, server_default=text("'running'"))
    error: Mapped[str | None] = mapped_column(Text, nullable=True)

    __table_args__ = (Index("ix_sync_log_host_started_at", "host", text("started_at DESC")),)


class SessionRow(BaseModel):
    """A session row in list results."""

    host: str
    tool: str
    session_id: str
    title: str
    working_directory: str
    model_name: str | None
    created_at: datetime | None
    last_activity_at: datetime | None
    message_count: int
    hidden: bool
    synced_at: datetime


class MessageRow(BaseModel):
    """A message row."""

    node_id: str
    parent_node_id: str | None
    message_id: str
    role: str
    content: str
    created_at: datetime | None
    metadata: dict[str, Any]


class SessionDetail(SessionRow):
    """A session with all messages loaded."""

    messages: list[MessageRow]
    metadata: dict[str, Any]


class MessageSearchResult(MessageRow):
    """A message search result with relevance ranking."""

    session_id: str
    session_title: str
    session_host: str
    session_tool: str
    session_working_directory: str
    rank: float  # PostgreSQL ts_rank score (FTS) or similarity score (semantic)


class SessionSearchResult(SessionRow):
    """A session search result with similarity score (semantic search)."""

    similarity: float  # cosine similarity (1.0 = identical, 0.0 = unrelated)
    matched_text: str  # excerpt of the text that matched


class HostStats(BaseModel):
    sessions: int
    messages: int
    last_sync: datetime | None


class ToolStats(BaseModel):
    sessions: int
    messages: int


class Stats(BaseModel):
    """Aggregate statistics."""

    total_sessions: int
    total_messages: int
    by_host: dict[str, HostStats]
    by_tool: dict[str, ToolStats]
    oldest_session: datetime | None
    newest_session: datetime | None


class HostInfo(BaseModel):
    host: str
    last_sync: datetime | None
    session_count: int
    message_count: int
