mcp/database/models.py

"""
Enterprise-Grade SQLAlchemy Database Models for CX AI Agent
"""
from datetime import datetime
from typing import Optional
from sqlalchemy import (
    Column, Integer, String, Text, DateTime, Float, Boolean,
    ForeignKey, Index, JSON, UniqueConstraint, CheckConstraint
)
from sqlalchemy.ext.asyncio import AsyncAttrs
from sqlalchemy.orm import DeclarativeBase, relationship, Mapped, mapped_column
from sqlalchemy.sql import func


class Base(AsyncAttrs, DeclarativeBase):
    """Base class for all models with async support"""
    pass


class TimestampMixin:
    """Mixin for created_at and updated_at timestamps"""
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
        nullable=False
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
        onupdate=func.now(),
        nullable=False
    )


class TenantMixin:
    """Mixin for multi-tenancy support"""
    tenant_id: Mapped[Optional[str]] = mapped_column(
        String(255),
        index=True,
        nullable=True,
        comment="Tenant ID for multi-tenancy isolation"
    )


class Company(Base, TimestampMixin, TenantMixin):
    """Company entity with rich metadata"""
    __tablename__ = "companies"

    id: Mapped[str] = mapped_column(String(255), primary_key=True)
    name: Mapped[str] = mapped_column(String(500), nullable=False, index=True)
    domain: Mapped[str] = mapped_column(String(500), nullable=False, unique=True, index=True)

    # Company details
    description: Mapped[Optional[str]] = mapped_column(Text)
    industry: Mapped[Optional[str]] = mapped_column(String(255), index=True)
    employee_count: Mapped[Optional[int]] = mapped_column(Integer)
    founded_year: Mapped[Optional[int]] = mapped_column(Integer)
    revenue_range: Mapped[Optional[str]] = mapped_column(String(100))
    funding: Mapped[Optional[str]] = mapped_column(String(255))

    # Location
    headquarters_city: Mapped[Optional[str]] = mapped_column(String(255))
    headquarters_state: Mapped[Optional[str]] = mapped_column(String(100))
    headquarters_country: Mapped[Optional[str]] = mapped_column(String(100), index=True)

    # Technology and social
    tech_stack: Mapped[Optional[dict]] = mapped_column(JSON)
    social_profiles: Mapped[Optional[dict]] = mapped_column(JSON)

    # Additional metadata
    metadata: Mapped[Optional[dict]] = mapped_column(JSON, default=dict)

    # Status
    is_active: Mapped[bool] = mapped_column(Boolean, default=True, index=True)

    # Relationships
    prospects: Mapped[list["Prospect"]] = relationship(
        "Prospect",
        back_populates="company",
        cascade="all, delete-orphan"
    )
    contacts: Mapped[list["Contact"]] = relationship(
        "Contact",
        back_populates="company",
        cascade="all, delete-orphan"
    )
    facts: Mapped[list["Fact"]] = relationship(
        "Fact",
        back_populates="company",
        cascade="all, delete-orphan"
    )

    __table_args__ = (
        Index('idx_company_domain_tenant', 'domain', 'tenant_id'),
        Index('idx_company_active_tenant', 'is_active', 'tenant_id'),
        Index('idx_company_industry_tenant', 'industry', 'tenant_id'),
    )

    def __repr__(self):
        return f"<Company(id={self.id}, name={self.name}, domain={self.domain})>"


class Prospect(Base, TimestampMixin, TenantMixin):
    """Prospect entity representing sales opportunities"""
    __tablename__ = "prospects"

    id: Mapped[str] = mapped_column(String(255), primary_key=True)
    company_id: Mapped[str] = mapped_column(
        String(255),
        ForeignKey("companies.id", ondelete="CASCADE"),
        nullable=False,
        index=True
    )

    # Scoring
    fit_score: Mapped[Optional[float]] = mapped_column(Float, index=True)
    engagement_score: Mapped[Optional[float]] = mapped_column(Float)
    intent_score: Mapped[Optional[float]] = mapped_column(Float)
    overall_score: Mapped[Optional[float]] = mapped_column(Float, index=True)

    # Status and stage
    status: Mapped[str] = mapped_column(
        String(50),
        default="new",
        index=True,
        comment="new, contacted, engaged, qualified, converted, lost"
    )
    stage: Mapped[str] = mapped_column(
        String(50),
        default="discovery",
        index=True,
        comment="discovery, qualification, proposal, negotiation, closed"
    )

    # Outreach tracking
    last_contacted_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
    last_replied_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
    emails_sent_count: Mapped[int] = mapped_column(Integer, default=0)
    emails_opened_count: Mapped[int] = mapped_column(Integer, default=0)
    emails_replied_count: Mapped[int] = mapped_column(Integer, default=0)

    # AI-generated content
    personalized_pitch: Mapped[Optional[str]] = mapped_column(Text)
    pain_points: Mapped[Optional[dict]] = mapped_column(JSON)
    value_propositions: Mapped[Optional[dict]] = mapped_column(JSON)

    # Metadata
    source: Mapped[Optional[str]] = mapped_column(String(255), comment="How was this prospect discovered")
    enrichment_data: Mapped[Optional[dict]] = mapped_column(JSON)
    metadata: Mapped[Optional[dict]] = mapped_column(JSON, default=dict)

    # Compliance
    is_suppressed: Mapped[bool] = mapped_column(Boolean, default=False, index=True)
    opt_out_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))

    # Relationships
    company: Mapped["Company"] = relationship("Company", back_populates="prospects")
    activities: Mapped[list["Activity"]] = relationship(
        "Activity",
        back_populates="prospect",
        cascade="all, delete-orphan",
        order_by="Activity.created_at.desc()"
    )
    handoffs: Mapped[list["Handoff"]] = relationship(
        "Handoff",
        back_populates="prospect",
        cascade="all, delete-orphan"
    )

    __table_args__ = (
        Index('idx_prospect_status_tenant', 'status', 'tenant_id'),
        Index('idx_prospect_stage_tenant', 'stage', 'tenant_id'),
        Index('idx_prospect_score_tenant', 'overall_score', 'tenant_id'),
        Index('idx_prospect_company_tenant', 'company_id', 'tenant_id'),
        CheckConstraint('fit_score >= 0 AND fit_score <= 100', name='check_fit_score_range'),
        CheckConstraint('overall_score >= 0 AND overall_score <= 100', name='check_overall_score_range'),
    )

    def __repr__(self):
        return f"<Prospect(id={self.id}, company_id={self.company_id}, score={self.overall_score})>"


class Contact(Base, TimestampMixin, TenantMixin):
    """Contact entity representing decision-makers"""
    __tablename__ = "contacts"

    id: Mapped[str] = mapped_column(String(255), primary_key=True)
    company_id: Mapped[str] = mapped_column(
        String(255),
        ForeignKey("companies.id", ondelete="CASCADE"),
        nullable=False,
        index=True
    )

    # Personal information
    email: Mapped[str] = mapped_column(String(500), nullable=False, unique=True, index=True)
    first_name: Mapped[Optional[str]] = mapped_column(String(255))
    last_name: Mapped[Optional[str]] = mapped_column(String(255))
    full_name: Mapped[Optional[str]] = mapped_column(String(500), index=True)

    # Professional information
    title: Mapped[Optional[str]] = mapped_column(String(500), index=True)
    department: Mapped[Optional[str]] = mapped_column(String(255), index=True)
    seniority: Mapped[Optional[str]] = mapped_column(
        String(50),
        comment="IC, Manager, Director, VP, C-Level"
    )

    # Contact details
    phone: Mapped[Optional[str]] = mapped_column(String(50))
    linkedin_url: Mapped[Optional[str]] = mapped_column(String(500))
    twitter_url: Mapped[Optional[str]] = mapped_column(String(500))

    # Validation
    email_valid: Mapped[bool] = mapped_column(Boolean, default=True, index=True)
    email_deliverability_score: Mapped[Optional[int]] = mapped_column(Integer)
    is_role_based: Mapped[bool] = mapped_column(Boolean, default=False, index=True)

    # Enrichment
    enrichment_data: Mapped[Optional[dict]] = mapped_column(JSON)
    metadata: Mapped[Optional[dict]] = mapped_column(JSON, default=dict)

    # Status
    is_active: Mapped[bool] = mapped_column(Boolean, default=True, index=True)
    is_primary_contact: Mapped[bool] = mapped_column(Boolean, default=False, index=True)

    # Relationships
    company: Mapped["Company"] = relationship("Company", back_populates="contacts")
    activities: Mapped[list["Activity"]] = relationship(
        "Activity",
        back_populates="contact",
        cascade="all, delete-orphan"
    )

    __table_args__ = (
        Index('idx_contact_email_tenant', 'email', 'tenant_id'),
        Index('idx_contact_company_tenant', 'company_id', 'tenant_id'),
        Index('idx_contact_valid_tenant', 'email_valid', 'tenant_id'),
        Index('idx_contact_seniority_tenant', 'seniority', 'tenant_id'),
    )

    def __repr__(self):
        return f"<Contact(id={self.id}, email={self.email}, title={self.title})>"


class Fact(Base, TimestampMixin, TenantMixin):
    """Fact entity for storing enrichment data and insights"""
    __tablename__ = "facts"

    id: Mapped[str] = mapped_column(String(255), primary_key=True)
    company_id: Mapped[str] = mapped_column(
        String(255),
        ForeignKey("companies.id", ondelete="CASCADE"),
        nullable=False,
        index=True
    )

    # Fact content
    fact_type: Mapped[str] = mapped_column(
        String(100),
        index=True,
        comment="news, funding, hiring, tech_stack, pain_point, etc."
    )
    title: Mapped[Optional[str]] = mapped_column(String(500))
    content: Mapped[str] = mapped_column(Text, nullable=False)

    # Source information
    source_url: Mapped[Optional[str]] = mapped_column(String(1000))
    source_name: Mapped[Optional[str]] = mapped_column(String(255))
    published_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), index=True)

    # Confidence and relevance
    confidence_score: Mapped[float] = mapped_column(Float, default=0.5)
    relevance_score: Mapped[Optional[float]] = mapped_column(Float)

    # Metadata
    metadata: Mapped[Optional[dict]] = mapped_column(JSON, default=dict)

    # Relationships
    company: Mapped["Company"] = relationship("Company", back_populates="facts")

    __table_args__ = (
        Index('idx_fact_company_tenant', 'company_id', 'tenant_id'),
        Index('idx_fact_type_tenant', 'fact_type', 'tenant_id'),
        Index('idx_fact_published_tenant', 'published_at', 'tenant_id'),
    )

    def __repr__(self):
        return f"<Fact(id={self.id}, type={self.fact_type}, company_id={self.company_id})>"


class Activity(Base, TimestampMixin, TenantMixin):
    """Activity entity for tracking all prospect interactions"""
    __tablename__ = "activities"

    id: Mapped[str] = mapped_column(String(255), primary_key=True)
    prospect_id: Mapped[str] = mapped_column(
        String(255),
        ForeignKey("prospects.id", ondelete="CASCADE"),
        nullable=False,
        index=True
    )
    contact_id: Mapped[Optional[str]] = mapped_column(
        String(255),
        ForeignKey("contacts.id", ondelete="SET NULL"),
        index=True
    )

    # Activity type
    activity_type: Mapped[str] = mapped_column(
        String(100),
        index=True,
        comment="email_sent, email_opened, email_replied, meeting_booked, call_made, etc."
    )
    direction: Mapped[str] = mapped_column(
        String(50),
        comment="inbound, outbound"
    )

    # Content
    subject: Mapped[Optional[str]] = mapped_column(String(1000))
    body: Mapped[Optional[str]] = mapped_column(Text)

    # Email specific
    email_thread_id: Mapped[Optional[str]] = mapped_column(String(255), index=True)
    email_message_id: Mapped[Optional[str]] = mapped_column(String(255))

    # Meeting specific
    meeting_scheduled_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), index=True)
    meeting_duration_minutes: Mapped[Optional[int]] = mapped_column(Integer)

    # Metadata
    metadata: Mapped[Optional[dict]] = mapped_column(JSON, default=dict)

    # Relationships
    prospect: Mapped["Prospect"] = relationship("Prospect", back_populates="activities")
    contact: Mapped[Optional["Contact"]] = relationship("Contact", back_populates="activities")

    __table_args__ = (
        Index('idx_activity_prospect_tenant', 'prospect_id', 'tenant_id'),
        Index('idx_activity_type_tenant', 'activity_type', 'tenant_id'),
        Index('idx_activity_thread_tenant', 'email_thread_id', 'tenant_id'),
        Index('idx_activity_created_tenant', 'created_at', 'tenant_id'),
    )

    def __repr__(self):
        return f"<Activity(id={self.id}, type={self.activity_type}, prospect_id={self.prospect_id})>"


class Suppression(Base, TimestampMixin, TenantMixin):
    """Suppression entity for compliance (opt-outs, bounces)"""
    __tablename__ = "suppressions"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

    # Suppression details
    suppression_type: Mapped[str] = mapped_column(
        String(50),
        index=True,
        comment="email, domain, opt_out, bounce, complaint"
    )
    value: Mapped[str] = mapped_column(String(500), nullable=False, index=True)

    # Reason
    reason: Mapped[Optional[str]] = mapped_column(String(500))
    source: Mapped[Optional[str]] = mapped_column(String(255))

    # Expiry
    expires_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), index=True)

    # Metadata
    metadata: Mapped[Optional[dict]] = mapped_column(JSON, default=dict)

    __table_args__ = (
        UniqueConstraint('suppression_type', 'value', 'tenant_id', name='uq_suppression_type_value_tenant'),
        Index('idx_suppression_type_value_tenant', 'suppression_type', 'value', 'tenant_id'),
        Index('idx_suppression_expires_tenant', 'expires_at', 'tenant_id'),
    )

    def __repr__(self):
        return f"<Suppression(type={self.suppression_type}, value={self.value})>"


class Handoff(Base, TimestampMixin, TenantMixin):
    """Handoff entity for AI-to-human sales transitions"""
    __tablename__ = "handoffs"

    id: Mapped[str] = mapped_column(String(255), primary_key=True)
    prospect_id: Mapped[str] = mapped_column(
        String(255),
        ForeignKey("prospects.id", ondelete="CASCADE"),
        nullable=False,
        index=True
    )

    # Handoff details
    status: Mapped[str] = mapped_column(
        String(50),
        default="pending",
        index=True,
        comment="pending, assigned, contacted, completed"
    )
    priority: Mapped[str] = mapped_column(
        String(50),
        default="medium",
        index=True,
        comment="low, medium, high, urgent"
    )

    # Assignment
    assigned_to: Mapped[Optional[str]] = mapped_column(String(255), index=True)
    assigned_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))

    # Summary
    summary: Mapped[Optional[str]] = mapped_column(Text)
    recommended_next_steps: Mapped[Optional[dict]] = mapped_column(JSON)
    conversation_history: Mapped[Optional[dict]] = mapped_column(JSON)

    # Metadata
    metadata: Mapped[Optional[dict]] = mapped_column(JSON, default=dict)

    # Relationships
    prospect: Mapped["Prospect"] = relationship("Prospect", back_populates="handoffs")

    __table_args__ = (
        Index('idx_handoff_prospect_tenant', 'prospect_id', 'tenant_id'),
        Index('idx_handoff_status_tenant', 'status', 'tenant_id'),
        Index('idx_handoff_assigned_tenant', 'assigned_to', 'tenant_id'),
    )

    def __repr__(self):
        return f"<Handoff(id={self.id}, prospect_id={self.prospect_id}, status={self.status})>"


class AuditLog(Base):
    """Audit log for compliance and security"""
    __tablename__ = "audit_logs"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

    # Who
    tenant_id: Mapped[Optional[str]] = mapped_column(String(255), index=True)
    user_id: Mapped[Optional[str]] = mapped_column(String(255), index=True)
    user_agent: Mapped[Optional[str]] = mapped_column(String(1000))
    ip_address: Mapped[Optional[str]] = mapped_column(String(50))

    # What
    action: Mapped[str] = mapped_column(String(100), nullable=False, index=True)
    resource_type: Mapped[str] = mapped_column(String(100), nullable=False, index=True)
    resource_id: Mapped[str] = mapped_column(String(255), nullable=False, index=True)

    # Changes
    old_value: Mapped[Optional[dict]] = mapped_column(JSON)
    new_value: Mapped[Optional[dict]] = mapped_column(JSON)

    # When
    timestamp: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
        nullable=False,
        index=True
    )

    # Additional context
    metadata: Mapped[Optional[dict]] = mapped_column(JSON, default=dict)

    __table_args__ = (
        Index('idx_audit_tenant_timestamp', 'tenant_id', 'timestamp'),
        Index('idx_audit_resource', 'resource_type', 'resource_id'),
        Index('idx_audit_action_timestamp', 'action', 'timestamp'),
    )

    def __repr__(self):
        return f"<AuditLog(id={self.id}, action={self.action}, resource={self.resource_type}/{self.resource_id})>"