from arepo.base import Base
from arepo.mixins import EntityLoaderMixin, AssociationLoaderMixin

from sqlalchemy.orm import relationship
from sqlalchemy import Column, Integer, String, ForeignKey, DateTime, ForeignKeyConstraint


class VulnerabilityModel(Base, EntityLoaderMixin):
    __tablename__ = "vulnerability"

    id = Column('id', String, primary_key=True)
    description = Column('description', String, nullable=True)
    published_date = Column('published_date', DateTime, nullable=False)
    last_modified_date = Column('last_modified_date', DateTime, nullable=False)
    configurations = relationship("ConfigurationModel", backref="vulnerability")
    source_id = Column(String, ForeignKey('source.id'), nullable=False)
    source = relationship("SourceModel", backref="vulnerabilities")

    # Relationship to ReferenceAssociationModel
    references = relationship(
        "ReferenceAssociationModel",
        back_populates="vulnerability"
    )

    cvss2 = relationship(
        "CVSS2AssociationModel",
        back_populates="vulnerability"
    )

    cvss3 = relationship(
        "CVSS3AssociationModel",
        back_populates="vulnerability"
    )


class VulnerabilityCWEModel(Base, AssociationLoaderMixin):
    __tablename__ = 'vulnerability_cwe'
    __table_args__ = (
        ForeignKeyConstraint(('vulnerability_id',), ['vulnerability.id']),
        ForeignKeyConstraint(('cwe_id',), ['cwe.id'])
    )

    vulnerability_id = Column(String, ForeignKey('vulnerability.id'), primary_key=True)
    cwe_id = Column(Integer, ForeignKey('cwe.id'), primary_key=True)
