from enum import IntEnum

class ClassificationLevel(IntEnum):
    """Data Classification Level - 7-Tier System"""

    PUBLIC = 1               # A1 - Public
    INTERNAL_PUBLIC = 2      # A2 - Internal Public
    INTERNAL_SENSITIVE = 3   # B1 - Internal Sensitive
    CONFIDENTIAL = 4         # B2 - Confidential
    HIGHLY_CONFIDENTIAL = 5  # C1 - Highly Confidential
    RESTRICTED = 6           # C2 - Restricted
    TOP_SECRET = 7           # D  - Top Secret

    @property
    def code(self) -> str:
        return {
            1: "A1", 2: "A2", 3: "B1", 4: "B2",
            5: "C1", 6: "C2", 7: "D",
        }[self.value]

    @property
    def requires_encryption_at_rest(self) -> bool:
        return self.value >= 4  # B2 and above

    @property
    def requires_audit_on_access(self) -> bool:
        return self.value >= 3  # B1 and above

    @property
    def max_retention_days(self) -> int | None:
        return {
            5: 365 * 3,   # C1: 3 years
            6: 365 * 5,   # C2: 5 years
            7: 365 * 10,  # D: 10 years
        }.get(self.value)

class ClassificationLabel(BaseModel):
    """Data Classification Label"""

    level: ClassificationLevel = Field(description="Classification level")
    category: DataCategory = Field(description="Data category")
    sub_category: str | None = Field(default=None, description="Sub-category")

    # Label metadata
    labeled_by: str = Field(description="Labeler (person/system)")
    labeled_at: datetime = Field(description="Labeling timestamp")
    review_status: ReviewStatus = Field(default=ReviewStatus.PENDING)
    reviewed_by: str | None = Field(default=None)
    reviewed_at: datetime | None = Field(default=None)

    # Label source
    source: LabelSource = Field(description="Label source")
    confidence: float = Field(default=1.0, ge=0.0, le=1.0, description="Confidence score")

    # Compliance tags
    regulatory_tags: list[str] = Field(
        default_factory=list,
        description="Associated regulatory tags (e.g., GDPR, HIPAA)",
    )


class DataCategory(str, Enum):
    PII = "pii"                    # Personally Identifiable Information
    PHI = "phi"                    # Protected Health Information
    PCI = "pci"                    # Payment Card Industry
    FINANCIAL = "financial"
    BUSINESS = "business"
    TECHNICAL = "technical"
    OPERATIONAL = "operational"


class LabelSource(str, Enum):
    MANUAL = "manual"
    RULE_ENGINE = "rule_engine"
    ML_CLASSIFIER = "ml_classifier"
    INHERITED = "inherited"
    SCHEMA_DEFAULT = "schema_default"

class PropertyDefinition(BaseModel):
    """Ontology property definition with classification"""

    name: str
    display_name: str
    data_type: DataType
    description: str = ""

    classification: ClassificationLabel | None = Field(
        default=None,
        description="Property data classification label",
    )

    min_classification: ClassificationLevel | None = Field(
        default=None,
        description="Minimum classification level allowed for this property",
    )

    field_security: FieldSecurityPolicy | None = Field(default=None)


class ObjectTypeDefinition(BaseModel):
    """Object type definition with object-level classification"""

    api_name: str
    display_name: str
    description: str

    effective_classification: ClassificationLevel = Field(
        default=ClassificationLevel.PUBLIC,
    )

    properties: dict[str, PropertyDefinition]

    def compute_effective_classification(self) -> ClassificationLevel:
        """Compute effective classification (highest across all properties)"""
        max_level = ClassificationLevel.PUBLIC
        for prop in self.properties.values():
            if prop.classification and prop.classification.level > max_level:
                max_level = prop.classification.level
        return max_level

class RuleBasedClassifier:
    """Rule-based automatic classification engine"""

    def __init__(self):
        self._rules: list[ClassificationRule] = []
        self._load_default_rules()

    def _load_default_rules(self):
        self._rules = [
            # PII rules
            ClassificationRule(
                name="pii_id_number",
                pattern=r"(id_card|ssn|social_security|national_id)",
                field_type_match=["string"],
                classification=ClassificationLevel.HIGHLY_CONFIDENTIAL,
                category=DataCategory.PII,
                confidence=0.95,
            ),
            ClassificationRule(
                name="pii_phone",
                pattern=r"(phone|mobile|tel|cell)",
                field_type_match=["string"],
                classification=ClassificationLevel.CONFIDENTIAL,
                category=DataCategory.PII,
                confidence=0.90,
            ),
            ClassificationRule(
                name="pii_email",
                pattern=r"(email|mail_address)",
                field_type_match=["string"],
                classification=ClassificationLevel.CONFIDENTIAL,
                category=DataCategory.PII,
                confidence=0.90,
            ),
            ClassificationRule(
                name="pii_name",
                pattern=r"(full_name|real_name|legal_name)",
                field_type_match=["string"],
                classification=ClassificationLevel.INTERNAL_SENSITIVE,
                category=DataCategory.PII,
                confidence=0.85,
            ),
            # Financial data rules
            ClassificationRule(
                name="financial_account",
                pattern=r"(bank_card|account_number|routing_number)",
                field_type_match=["string"],
                classification=ClassificationLevel.HIGHLY_CONFIDENTIAL,
                category=DataCategory.PCI,
                confidence=0.95,
            ),
            ClassificationRule(
                name="financial_salary",
                pattern=r"(salary|compensation|wage|pay_rate)",
                field_type_match=["decimal", "float", "integer"],
                classification=ClassificationLevel.CONFIDENTIAL,
                category=DataCategory.FINANCIAL,
                confidence=0.90,
            ),
            # Health data rules
            ClassificationRule(
                name="phi_diagnosis",
                pattern=r"(diagnosis|medical_record|health_condition)",
                field_type_match=["string", "text"],
                classification=ClassificationLevel.HIGHLY_CONFIDENTIAL,
                category=DataCategory.PHI,
                confidence=0.95,
            ),
        ]

    def classify_field(
        self,
        field_name: str,
        field_type: str,
        description: str = "",
    ) -> ClassificationLabel | None:
        best_match: ClassificationRule | None = None
        best_confidence = 0.0

        text_to_match = f"{field_name} {description}".lower()

        for rule in self._rules:
            if field_type not in rule.field_type_match:
                continue
            if re.search(rule.pattern, text_to_match, re.IGNORECASE):
                if rule.confidence > best_confidence:
                    best_match = rule
                    best_confidence = rule.confidence

        if best_match is None:
            return None

        return ClassificationLabel(
            level=best_match.classification,
            category=best_match.category,
            labeled_by="rule_engine",
            labeled_at=datetime.utcnow(),
            source=LabelSource.RULE_ENGINE,
            confidence=best_match.confidence,
        )

class ClassificationPropagator:
    """Classification label propagator"""

    class PropagationPolicy(str, Enum):
        HIGH_WATER_MARK = "high_water_mark"
        EXPLICIT_ONLY = "explicit_only"
        INHERIT_SOURCE = "inherit_source"

    def propagate_on_transform(
        self,
        source_classifications: list[ClassificationLabel],
        transform_type: str,
        policy: PropagationPolicy = PropagationPolicy.HIGH_WATER_MARK,
    ) -> ClassificationLabel:
        if policy == self.PropagationPolicy.HIGH_WATER_MARK:
            max_level = max(c.level for c in source_classifications)
            categories = set(c.category for c in source_classifications)
            regulatory_tags = set()
            for c in source_classifications:
                regulatory_tags.update(c.regulatory_tags)

            return ClassificationLabel(
                level=max_level,
                category=categories.pop() if len(categories) == 1 else DataCategory.BUSINESS,
                labeled_by="propagation_engine",
                labeled_at=datetime.utcnow(),
                source=LabelSource.INHERITED,
                confidence=min(c.confidence for c in source_classifications),
                regulatory_tags=list(regulatory_tags),
            )

        elif policy == self.PropagationPolicy.EXPLICIT_ONLY:
            raise ValueError("Explicit policy requires manual classification")

        elif policy == self.PropagationPolicy.INHERIT_SOURCE:
            return source_classifications[0].model_copy(
                update={
                    "source": LabelSource.INHERITED,
                    "labeled_at": datetime.utcnow(),
                }
            )

syntax = "proto3";

package onto.classification.v1;

enum ClassificationLevel {
    CLASSIFICATION_LEVEL_UNSPECIFIED = 0;
    PUBLIC = 1;
    INTERNAL_PUBLIC = 2;
    INTERNAL_SENSITIVE = 3;
    CONFIDENTIAL = 4;
    HIGHLY_CONFIDENTIAL = 5;
    RESTRICTED = 6;
    TOP_SECRET = 7;
}

enum DataCategory {
    DATA_CATEGORY_UNSPECIFIED = 0;
    PII = 1;
    PHI = 2;
    PCI = 3;
    FINANCIAL = 4;
    BUSINESS = 5;
    TECHNICAL = 6;
    OPERATIONAL = 7;
}

message ClassificationLabel {
    ClassificationLevel level = 1;
    DataCategory category = 2;
    string sub_category = 3;
    string labeled_by = 4;
    google.protobuf.Timestamp labeled_at = 5;
    string source = 6;
    double confidence = 7;
    repeated string regulatory_tags = 8;
}

message FieldClassification {
    string field_name = 1;
    ClassificationLabel label = 2;
}

message ObjectClassification {
    string object_type = 1;
    ClassificationLevel effective_level = 2;
    repeated FieldClassification field_classifications = 3;
}

class ClassificationInterceptor(grpc.aio.UnaryUnaryClientInterceptor):
    """gRPC client interceptor - propagate classification context"""

    async def intercept_unary_unary(self, continuation, client_call_details, request):
        current_classification = classification_context.get()

        if current_classification:
            metadata = list(client_call_details.metadata or [])
            metadata.append((
                "x-classification-level",
                str(current_classification.level.value),
            ))
            metadata.append((
                "x-classification-category",
                current_classification.category.value,
            ))

            new_details = client_call_details._replace(metadata=metadata)
            return await continuation(new_details, request)

        return await continuation(client_call_details, request)

class ClassificationBasedAccessControl:
    """Classification-based access control"""

    CLEARANCE_MAP = {
        "public_user": ClassificationLevel.PUBLIC,
        "internal_user": ClassificationLevel.INTERNAL_SENSITIVE,
        "confidential_user": ClassificationLevel.CONFIDENTIAL,
        "secret_user": ClassificationLevel.HIGHLY_CONFIDENTIAL,
        "top_secret_user": ClassificationLevel.TOP_SECRET,
    }

    def check_access(
        self,
        user_clearance: str,
        resource_classification: ClassificationLevel,
    ) -> AccessDecision:
        max_allowed = self.CLEARANCE_MAP.get(user_clearance, ClassificationLevel.PUBLIC)

        if resource_classification <= max_allowed:
            return AccessDecision.ALLOW
        else:
            return AccessDecision.DENY

CLASSIFICATION_MASKING_MAP = {
    ClassificationLevel.PUBLIC: MaskingMode.NONE,
    ClassificationLevel.INTERNAL_PUBLIC: MaskingMode.NONE,
    ClassificationLevel.INTERNAL_SENSITIVE: MaskingMode.PARTIAL,
    ClassificationLevel.CONFIDENTIAL: MaskingMode.PARTIAL,
    ClassificationLevel.HIGHLY_CONFIDENTIAL: MaskingMode.HASH,
    ClassificationLevel.RESTRICTED: MaskingMode.FULL,
    ClassificationLevel.TOP_SECRET: MaskingMode.FULL,
}

class ClassificationAuditTrigger:
    """Classification-driven audit trigger"""

    AUDIT_THRESHOLDS = {
        ClassificationLevel.INTERNAL_SENSITIVE: AuditLevel.BASIC,
        ClassificationLevel.CONFIDENTIAL: AuditLevel.STANDARD,
        ClassificationLevel.HIGHLY_CONFIDENTIAL: AuditLevel.DETAILED,
        ClassificationLevel.RESTRICTED: AuditLevel.COMPREHENSIVE,
        ClassificationLevel.TOP_SECRET: AuditLevel.FORENSIC,
    }

    async def on_data_access(
        self,
        classification: ClassificationLevel,
        access_event: DataAccessEvent,
    ) -> None:
        audit_level = self.AUDIT_THRESHOLDS.get(classification)
        if audit_level:
            await self._audit_service.record(
                event=access_event,
                level=audit_level,
                classification=classification,
            )

Schema Registration → Auto-Classification Engine → Human Review → Label Activation
         │                     │                       │              │
         │               Rule matching          Security team    ABAC policies
         │               ML classifier         Compliance team   Masking policies
         │                                                        Audit policies

class ClassificationChangeRequest(BaseModel):
    """Classification change request"""

    id: str
    object_type: str
    field_name: str
    current_classification: ClassificationLevel
    requested_classification: ClassificationLevel
    reason: str
    requester: str
    requested_at: datetime

    approvers: list[str] = Field(default_factory=list)
    approval_status: ApprovalStatus = Field(default=ApprovalStatus.PENDING)

    @validator("requested_classification")
    def validate_downgrade(cls, v, values):
        """Downgrades require additional approval"""
        current = values.get("current_classification")
        if current and v < current:
            values["approvers"] = ["security_team_lead", "compliance_officer"]
        return v

class ClassificationReviewScheduler:
    """Classification review scheduler"""

    REVIEW_INTERVALS = {
        ClassificationLevel.TOP_SECRET: timedelta(days=90),
        ClassificationLevel.RESTRICTED: timedelta(days=180),
        ClassificationLevel.HIGHLY_CONFIDENTIAL: timedelta(days=365),
        ClassificationLevel.CONFIDENTIAL: timedelta(days=365 * 2),
    }

    async def schedule_reviews(self) -> list[ReviewTask]:
        tasks = []
        for obj_type in await self._schema_registry.list_object_types():
            for field_name, prop in obj_type.properties.items():
                if prop.classification is None:
                    continue
                interval = self.REVIEW_INTERVALS.get(prop.classification.level)
                if interval is None:
                    continue
                if datetime.utcnow() - prop.classification.labeled_at > interval:
                    tasks.append(ReviewTask(
                        object_type=obj_type.api_name,
                        field_name=field_name,
                        current_classification=prop.classification,
                        due_date=prop.classification.labeled_at + interval,
                    ))
        return tasks

class TestClassification:
    def test_auto_classify_pii_fields(self):
        classifier = RuleBasedClassifier()
        result = classifier.classify_field("ssn", "string", "Social Security Number")
        assert result is not None
        assert result.level == ClassificationLevel.HIGHLY_CONFIDENTIAL
        assert result.category == DataCategory.PII

    def test_propagation_high_water_mark(self):
        propagator = ClassificationPropagator()
        sources = [
            make_label(ClassificationLevel.INTERNAL_SENSITIVE),
            make_label(ClassificationLevel.CONFIDENTIAL),
        ]
        result = propagator.propagate_on_transform(sources, "join")
        assert result.level == ClassificationLevel.CONFIDENTIAL

    def test_object_effective_classification(self):
        obj = ObjectTypeDefinition(
            api_name="Customer",
            display_name="Customer",
            description="Customer object",
            properties={
                "name": PropertyDefinition(
                    name="name", display_name="Name", data_type=DataType.STRING,
                    classification=make_label(ClassificationLevel.INTERNAL_SENSITIVE),
                ),
                "ssn": PropertyDefinition(
                    name="ssn", display_name="SSN", data_type=DataType.STRING,
                    classification=make_label(ClassificationLevel.HIGHLY_CONFIDENTIAL),
                ),
            },
        )
        assert obj.compute_effective_classification() == ClassificationLevel.HIGHLY_CONFIDENTIAL