Source code for test_data_workbench.core.models

"""Core data models for V4 adaptive framework."""

from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, Any, Set
import uuid




class EntityType(Enum):
    """Common e-commerce entity types for classification."""
    USER = "user"
    CUSTOMER = "customer"
    PRODUCT = "product"
    ORDER = "order"
    REVIEW = "review"
    CATEGORY = "category"
    INVENTORY = "inventory"
    PAYMENT = "payment"
    SHIPPING = "shipping"
    ANALYTICS = "analytics"
    UNKNOWN = "unknown"





class ConstraintType(Enum):
    """Database constraint types."""
    PRIMARY_KEY = "primary_key"
    FOREIGN_KEY = "foreign_key"
    UNIQUE = "unique"
    NOT_NULL = "not_null"
    CHECK = "check"
    INDEX = "index"





@dataclass
class Column:
    """Database column metadata."""
    name: str
    data_type: str
    nullable: bool = True
    default: Optional[str] = None
    max_length: Optional[int] = None
    constraints: List[ConstraintType] = field(default_factory=list)
    sample_values: List[Any] = field(default_factory=list)
    
    @property
    def is_primary_key(self) -> bool:
        return ConstraintType.PRIMARY_KEY in self.constraints
    
    @property
    def is_foreign_key(self) -> bool:
        return ConstraintType.FOREIGN_KEY in self.constraints





@dataclass
class Relationship:
    """Table relationship metadata."""
    from_table: str
    to_table: str
    from_column: str
    to_column: str
    relationship_type: str  # "one_to_one", "one_to_many", "many_to_many"
    strength: float = 1.0  # Confidence score 0-1





@dataclass
class Table:
    """Database table metadata."""
    name: str
    columns: List[Column]
    row_count: Optional[int] = 0
    entity_type: EntityType = EntityType.UNKNOWN
    relationships: List[Relationship] = field(default_factory=list)
    
    @property
    def primary_keys(self) -> List[Column]:
        return [col for col in self.columns if col.is_primary_key]
    
    @property
    def foreign_keys(self) -> List[Column]:
        return [col for col in self.columns if col.is_foreign_key]





@dataclass
class BusinessRule:
    """Inferred business rule from data analysis."""
    rule_id: str
    description: str
    rule_type: str  # "format", "range", "dependency", "pattern"
    affected_columns: List[str]
    confidence: float  # 0-1 confidence score
    examples: List[str] = field(default_factory=list)
    


    @classmethod
    def create(cls, description: str, rule_type: str, columns: List[str], 
               confidence: float = 0.8) -> 'BusinessRule':
        return cls(
            rule_id=str(uuid.uuid4())[:8],
            description=description,
            rule_type=rule_type,
            affected_columns=columns,
            confidence=confidence
        )






@dataclass
class PIIColumn:
    """A column flagged as likely containing personally identifiable
    information, based on a name-based heuristic (no row data needed)."""
    table: str
    column: str
    category: str





@dataclass
class SchemaInfo:
    """Complete database schema analysis results."""
    database_name: str
    tables: List[Table]
    relationships: List[Relationship]
    business_rules: List[BusinessRule]
    analysis_metadata: Dict[str, Any] = field(default_factory=dict)
    pii_columns: List[PIIColumn] = field(default_factory=list)
    
    @property
    def entity_types(self) -> Dict[EntityType, List[Table]]:
        """Group tables by detected entity type."""
        groups = {}
        for table in self.tables:
            if table.entity_type not in groups:
                groups[table.entity_type] = []
            groups[table.entity_type].append(table)
        return groups
    
    @property
    def table_names(self) -> Set[str]:
        return {table.name for table in self.tables}
    


    def get_table(self, name: str) -> Optional[Table]:
        """Get table by name."""
        for table in self.tables:
            if table.name == name:
                return table
        return None






@dataclass
class GeneratorConfig:
    """Configuration for a specific entity generator."""
    entity_type: EntityType
    table_name: str
    dependencies: List[str] = field(default_factory=list)
    generation_rules: Dict[str, Any] = field(default_factory=dict)
    fallback_config: Optional[Dict[str, Any]] = None





@dataclass
class ScenarioConfig:
    """Business scenario configuration."""
    name: str
    description: str
    entities: Dict[str, int]  # entity_type -> count
    relationships: List[str] = field(default_factory=list)
    constraints: Dict[str, Any] = field(default_factory=dict)





@dataclass
class AdaptationResult:
    """Results from rapid adaptation workflow."""
    schema_info: SchemaInfo
    generated_code: Dict[str, str]  # filename -> code content
    config_files: Dict[str, str]    # filename -> config content
    templates: Dict[str, str]       # template_type -> template content
    analysis_time: float
    success: bool
    errors: List[str] = field(default_factory=list)
    # One-line summary from the post-generation constraint check (see
    # adaptation.constraint_check), e.g. "Constraint check: 12 checks
    # passed". None when the check couldn't run (e.g. empty schema) or
    # was skipped after a best-effort failure.
    constraint_check_summary: Optional[str] = None





@dataclass
class ValidationResult:
    """Results from team contribution validation."""
    valid: bool
    errors: List[str] = field(default_factory=list)
    warnings: List[str] = field(default_factory=list)
    suggestions: List[str] = field(default_factory=list)





@dataclass
class DemoScenario:
    """Demonstration scenario configuration."""
    name: str
    description: str
    required_features: List[str]
    fallback_features: List[str]
    demo_data: Dict[str, Any]
    success_criteria: List[str]





@dataclass
class DemoStatus:
    """Current demonstration readiness status."""
    ready_scenarios: List[DemoScenario]
    blocked_scenarios: List[DemoScenario] 
    fallback_plan: Optional[DemoScenario]
    confidence_score: float