"""Database models for tracking server using SQLAlchemy (SQLite or Postgres)."""

from typing import Any, Dict, Type

from sqlalchemy import Column, ForeignKey, Index, Integer, String, Text, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session, relationship, sessionmaker

Base: Type[Any] = declarative_base()


class Tag(Base):  # type: ignore[misc]
    """Tag stores key-value metadata for runs (git info, environment, user, etc.).

    Similar to MLflow's tag system - makes metadata easily searchable.
    """

    __tablename__ = "tags"

    id = Column(Integer, primary_key=True, autoincrement=True)
    run_id = Column(String, ForeignKey("runs.run_id"), nullable=False)
    key = Column(String, nullable=False)  # e.g., "git.commit", "user", "git.branch"
    value = Column(String, nullable=False)  # e.g., "abc123", "brandon", "main"

    # Relationship
    run = relationship("Run", back_populates="tags")

    # Indexes for efficient queries
    __table_args__ = (
        Index("idx_tags_run_id", "run_id"),
        Index("idx_tags_key", "key"),
        Index("idx_tags_key_value", "key", "value"),
        # Ensure unique key per run (can't have duplicate keys for same run)
        Index("idx_tags_unique", "run_id", "key", unique=True),
    )


class Run(Base):  # type: ignore[misc]
    """Run represents a single training run.

    Similar to MLflow Run, contains metadata and lifecycle info.
    """

    __tablename__ = "runs"

    # Primary key - generated by artifacta SDK
    run_id = Column(String, primary_key=True)

    # Auto-generated display name (e.g., "Run 1", "Run 2")
    # This is what the UI displays - solves the "trial_1" collision issue
    name = Column(String, nullable=False, unique=True)

    # Project/experiment grouping (optional)
    project = Column(String, nullable=True)

    # Link to config artifact (single source of truth!)
    config_artifact_id = Column(String, ForeignKey("artifacts.artifact_id"), nullable=True)

    # Lifecycle timestamps
    created_at = Column(Integer, nullable=False)  # Unix timestamp (ms)

    # Relationships
    tags = relationship("Tag", back_populates="run", cascade="all, delete-orphan")

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for API responses."""
        return {
            "run_id": self.run_id,
            "name": self.name,
            "project": self.project,
            "config_artifact_id": self.config_artifact_id,
            "created_at": self.created_at,
        }


class StructuredData(Base):  # type: ignore[misc]
    """Stores primitives (Series, Distribution, etc.)."""

    __tablename__ = "structured_data"

    id = Column(Integer, primary_key=True, autoincrement=True)
    run_id = Column(String, ForeignKey("runs.run_id"), nullable=False)
    name = Column(String, nullable=False)
    primitive_type = Column(String, nullable=False)
    section = Column(String, nullable=True)
    data = Column(Text, nullable=False)
    meta = Column(Text, nullable=True)
    timestamp = Column(Integer, nullable=False)

    __table_args__ = (
        Index("idx_structured_data_run_id", "run_id"),
        Index("idx_structured_data_type", "primitive_type"),
        Index("idx_structured_data_name", "name"),
    )


class Artifact(Base):  # type: ignore[misc]
    """Stores artifact metadata and content."""

    __tablename__ = "artifacts"

    artifact_id = Column(String, primary_key=True)
    run_id = Column(String, ForeignKey("runs.run_id"), nullable=False)
    name = Column(String, nullable=False)
    hash = Column(String, nullable=False)
    storage_path = Column(String, nullable=False)
    size_bytes = Column(Integer, nullable=True)
    meta = Column(Text, nullable=True)
    content = Column(Text, nullable=True)
    created_at = Column(Integer, nullable=False)

    __table_args__ = (
        Index("idx_artifacts_run_id", "run_id"),
        Index("idx_artifacts_hash", "hash"),
    )


class ArtifactLink(Base):  # type: ignore[misc]
    """Links artifacts to runs with role (input/output)."""

    __tablename__ = "artifact_links"

    link_id = Column(String, primary_key=True)
    artifact_id = Column(String, ForeignKey("artifacts.artifact_id"), nullable=False)
    run_id = Column(String, ForeignKey("runs.run_id"), nullable=False)
    role = Column(String, nullable=False)  # 'input' or 'output'
    created_at = Column(Integer, nullable=False)

    __table_args__ = (
        Index("idx_artifact_links_run_id", "run_id"),
        Index("idx_artifact_links_artifact_id", "artifact_id"),
    )


class Project(Base):  # type: ignore[misc]
    """Project represents a collection of related runs.

    Auto-created when a run specifies a project name.
    """

    __tablename__ = "projects"

    project_id = Column(String, primary_key=True)
    created_at = Column(Integer, nullable=False)
    updated_at = Column(Integer, nullable=False)

    # Relationships
    notes = relationship("ProjectNote", back_populates="project", cascade="all, delete-orphan")

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for API responses."""
        return {
            "project_id": self.project_id,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
        }


class ProjectNote(Base):  # type: ignore[misc]
    """ProjectNote represents a lab notebook entry for a project.

    Contains markdown content and can link to multiple runs.
    """

    __tablename__ = "project_notes"

    id = Column(Integer, primary_key=True, autoincrement=True)
    project_id = Column(String, ForeignKey("projects.project_id"), nullable=False)
    title = Column(String, nullable=False)
    content = Column(Text, nullable=False)  # Markdown content
    created_at = Column(Integer, nullable=False)
    updated_at = Column(Integer, nullable=False)

    # Relationships
    project = relationship("Project", back_populates="notes")
    attachments = relationship(
        "ProjectNoteAttachment", back_populates="note", cascade="all, delete-orphan"
    )

    # Indexes
    __table_args__ = (
        Index("idx_project_notes_project_id", "project_id"),
        Index("idx_project_notes_created_at", "created_at"),
    )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for API responses."""
        return {
            "id": self.id,
            "project_id": self.project_id,
            "title": self.title,
            "content": self.content,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "attachments": [att.to_dict() for att in self.attachments],
        }


class ProjectNoteAttachment(Base):  # type: ignore[misc]
    """ProjectNoteAttachment represents a file attached to a project note.

    Files are stored with hash-based paths for integrity and deduplication.
    """

    __tablename__ = "project_note_attachments"

    id = Column(Integer, primary_key=True, autoincrement=True)
    note_id = Column(Integer, ForeignKey("project_notes.id", ondelete="CASCADE"), nullable=False)
    real_name = Column(String, nullable=False)  # Original filename
    storage_path = Column(String, nullable=False)  # Hash-based path: uploads/ab/ab34...ext
    mime_type = Column(String, nullable=False)
    filesize = Column(Integer, nullable=False)
    hash = Column(String, nullable=True)  # SHA256 hash
    created_at = Column(Integer, nullable=False)

    # Relationships
    note = relationship("ProjectNote", back_populates="attachments")

    # Indexes
    __table_args__ = (Index("idx_attachments_note_id", "note_id"),)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for API responses."""
        return {
            "id": self.id,
            "note_id": self.note_id,
            "real_name": self.real_name,
            "storage_path": self.storage_path,
            "mime_type": self.mime_type,
            "filesize": self.filesize,
            "hash": self.hash,
            "created_at": self.created_at,
        }


class Database:
    """Database manager - handles connection and session management.

    Similar to MLflow's SqlAlchemyStore.
    """

    def __init__(self, db_uri: str = "sqlite:///./data/runs.db"):
        """Initialize database connection.

        Args:
            db_uri: SQLAlchemy database URI
                    - SQLite: "sqlite:///./data/runs.db"
                    - Postgres: "postgresql://user:pass@host:port/dbname"
        """
        self.engine = create_engine(
            db_uri,
            # SQLite-specific settings
            connect_args={"check_same_thread": False} if db_uri.startswith("sqlite") else {},
        )
        self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)

        # Create tables if they don't exist
        Base.metadata.create_all(bind=self.engine)

    def get_session(self) -> Session:
        """Get a new database session."""
        return self.SessionLocal()

    def generate_run_name(self, run_id: str) -> str:
        """Auto-generate unique run name from run_id using hash.

        Uses first 8 chars of SHA256 hash for compact, collision-resistant names
        Format: "run_XXXXXXXX" (e.g., "run_a3f4b2c1")

        Guaranteed unique since run_id is unique, no assumptions about format
        """
        import hashlib

        hash_digest = hashlib.sha256(run_id.encode()).hexdigest()
        return f"run_{hash_digest[:8]}"


def init_db() -> None:
    """Initialize the database (for CLI usage)."""
    import os
    from pathlib import Path

    # Get project root (parent of tracking-server directory)
    project_root = Path(__file__).parent.parent
    db_path = os.getenv("DATABASE_PATH", str(project_root / "data" / "runs.db"))
    db_uri = f"sqlite:///{db_path}"

    # Ensure directory exists
    Path(db_path).parent.mkdir(parents=True, exist_ok=True)

    # Create database instance (this will create tables)
    Database(db_uri=db_uri)
