from datetime import datetime
from typing import Optional
from sqlalchemy import (
    Column, String, Text, Boolean, Integer, BigInteger,
    DateTime, ForeignKey, CheckConstraint, Index, JSON
)
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import relationship, declarative_base
import uuid

Base = declarative_base()


class Grove(Base):
    __tablename__ = "groves"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    name = Column(String(255), unique=True, nullable=False)
    description = Column(Text)
    is_public = Column(Boolean, default=True)
    created_at = Column(DateTime(timezone=True), default=datetime.utcnow)
    updated_at = Column(DateTime(timezone=True), default=datetime.utcnow, onupdate=datetime.utcnow)
    created_by = Column(String(255), nullable=False)

    trees = relationship("Tree", back_populates="grove", cascade="all, delete-orphan")
    permissions = relationship("AccessPermission", back_populates="grove", cascade="all, delete-orphan")

    __table_args__ = (
        Index("idx_groves_name", "name"),
        Index("idx_groves_created_by", "created_by"),
    )


class Tree(Base):
    __tablename__ = "trees"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    grove_id = Column(UUID(as_uuid=True), ForeignKey("groves.id", ondelete="CASCADE"), nullable=False)
    name = Column(String(255), nullable=False)
    description = Column(Text)
    created_at = Column(DateTime(timezone=True), default=datetime.utcnow)
    updated_at = Column(DateTime(timezone=True), default=datetime.utcnow, onupdate=datetime.utcnow)

    grove = relationship("Grove", back_populates="trees")
    grafts = relationship("Graft", back_populates="tree", cascade="all, delete-orphan")
    harvests = relationship("Harvest", back_populates="tree", cascade="all, delete-orphan")
    consumers = relationship("Consumer", back_populates="tree", cascade="all, delete-orphan")

    __table_args__ = (
        Index("idx_trees_grove_id", "grove_id"),
        Index("idx_trees_name", "name"),
        {"extend_existing": True},
    )


class Fruit(Base):
    __tablename__ = "fruits"

    id = Column(String(64), primary_key=True)  # SHA256 hash
    size = Column(BigInteger, nullable=False)
    content_type = Column(String(255))
    original_name = Column(String(1024))
    created_at = Column(DateTime(timezone=True), default=datetime.utcnow)
    created_by = Column(String(255), nullable=False)
    ref_count = Column(Integer, default=1)
    s3_key = Column(String(1024), nullable=False)

    grafts = relationship("Graft", back_populates="fruit")
    harvests = relationship("Harvest", back_populates="fruit")

    __table_args__ = (
        Index("idx_fruits_created_at", "created_at"),
        Index("idx_fruits_created_by", "created_by"),
    )


class Graft(Base):
    __tablename__ = "grafts"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    tree_id = Column(UUID(as_uuid=True), ForeignKey("trees.id", ondelete="CASCADE"), nullable=False)
    name = Column(String(255), nullable=False)
    fruit_id = Column(String(64), ForeignKey("fruits.id"), nullable=False)
    created_at = Column(DateTime(timezone=True), default=datetime.utcnow)
    created_by = Column(String(255), nullable=False)

    tree = relationship("Tree", back_populates="grafts")
    fruit = relationship("Fruit", back_populates="grafts")
    history = relationship("GraftHistory", back_populates="graft", cascade="all, delete-orphan")

    __table_args__ = (
        Index("idx_grafts_tree_id", "tree_id"),
        Index("idx_grafts_fruit_id", "fruit_id"),
    )


class GraftHistory(Base):
    __tablename__ = "graft_history"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    graft_id = Column(UUID(as_uuid=True), ForeignKey("grafts.id", ondelete="CASCADE"), nullable=False)
    old_fruit_id = Column(String(64), ForeignKey("fruits.id"))
    new_fruit_id = Column(String(64), ForeignKey("fruits.id"), nullable=False)
    changed_at = Column(DateTime(timezone=True), default=datetime.utcnow)
    changed_by = Column(String(255), nullable=False)

    graft = relationship("Graft", back_populates="history")

    __table_args__ = (
        Index("idx_graft_history_graft_id", "graft_id"),
    )


class Harvest(Base):
    __tablename__ = "harvests"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    fruit_id = Column(String(64), ForeignKey("fruits.id"), nullable=False)
    tree_id = Column(UUID(as_uuid=True), ForeignKey("trees.id"), nullable=False)
    original_name = Column(String(1024))
    harvested_at = Column(DateTime(timezone=True), default=datetime.utcnow)
    harvested_by = Column(String(255), nullable=False)
    source_ip = Column(String(45))

    fruit = relationship("Fruit", back_populates="harvests")
    tree = relationship("Tree", back_populates="harvests")

    __table_args__ = (
        Index("idx_harvests_fruit_id", "fruit_id"),
        Index("idx_harvests_tree_id", "tree_id"),
        Index("idx_harvests_harvested_at", "harvested_at"),
    )


class Consumer(Base):
    __tablename__ = "consumers"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    tree_id = Column(UUID(as_uuid=True), ForeignKey("trees.id", ondelete="CASCADE"), nullable=False)
    project_url = Column(String(2048), nullable=False)
    last_access = Column(DateTime(timezone=True), default=datetime.utcnow)
    created_at = Column(DateTime(timezone=True), default=datetime.utcnow)

    tree = relationship("Tree", back_populates="consumers")

    __table_args__ = (
        Index("idx_consumers_tree_id", "tree_id"),
        Index("idx_consumers_last_access", "last_access"),
    )


class AccessPermission(Base):
    __tablename__ = "access_permissions"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    grove_id = Column(UUID(as_uuid=True), ForeignKey("groves.id", ondelete="CASCADE"), nullable=False)
    user_id = Column(String(255), nullable=False)
    level = Column(String(20), nullable=False)
    created_at = Column(DateTime(timezone=True), default=datetime.utcnow)
    expires_at = Column(DateTime(timezone=True))

    grove = relationship("Grove", back_populates="permissions")

    __table_args__ = (
        CheckConstraint("level IN ('read', 'write', 'admin')", name="check_level"),
        Index("idx_access_permissions_grove_id", "grove_id"),
        Index("idx_access_permissions_user_id", "user_id"),
    )


class APIKey(Base):
    __tablename__ = "api_keys"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    key_hash = Column(String(64), unique=True, nullable=False)
    name = Column(String(255), nullable=False)
    user_id = Column(String(255), nullable=False)
    created_at = Column(DateTime(timezone=True), default=datetime.utcnow)
    expires_at = Column(DateTime(timezone=True))
    last_used = Column(DateTime(timezone=True))

    __table_args__ = (
        Index("idx_api_keys_user_id", "user_id"),
        Index("idx_api_keys_key_hash", "key_hash"),
    )


class AuditLog(Base):
    __tablename__ = "audit_logs"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    action = Column(String(100), nullable=False)
    resource = Column(String(1024), nullable=False)
    user_id = Column(String(255), nullable=False)
    details = Column(JSON)
    timestamp = Column(DateTime(timezone=True), default=datetime.utcnow)
    source_ip = Column(String(45))

    __table_args__ = (
        Index("idx_audit_logs_action", "action"),
        Index("idx_audit_logs_resource", "resource"),
        Index("idx_audit_logs_user_id", "user_id"),
        Index("idx_audit_logs_timestamp", "timestamp"),
    )