"""
eaigins_sdk.classification
──────────────────────────
Data classification system for AI governance.

Answers question 6 (Security) — organizations can define data classification
policies that automatically govern AI access to sensitive data sources.

Classification levels (aligned with enterprise data governance standards):
    PUBLIC      — freely accessible, no governance required
    INTERNAL    — internal use, basic logging required
    CONFIDENTIAL — PII, financial, requires approval for bulk access
    RESTRICTED  — PHI, payment card data, requires human authorization
    CRITICAL    — psychiatric records, biometric data, always blocked for AI

Usage::

    from eaigins_sdk.classification import (
        DataClassifier, DataClass,
        classify_context_sources, build_classification_policies,
    )
    from eaigins_sdk import ControlPlane
    from eaigins_sdk.policies import AML_POLICIES

    # 1. Define your org's data classification scheme
    classifier = DataClassifier()

    # Override or extend built-in classifications
    classifier.register("my_proprietary_db",     DataClass.RESTRICTED)
    classifier.register("public_press_releases",  DataClass.PUBLIC)
    classifier.register("employee_salary_data",   DataClass.CRITICAL)

    # 2. Auto-generate policies from classification rules
    classification_policies = build_classification_policies(
        classifier,
        restricted_requires_approval = True,
        critical_always_blocks       = True,
    )

    # 3. Combine with domain policies
    cp = ControlPlane(policies=AML_POLICIES + classification_policies)

    # 4. Check classification of a data source
    level = classifier.classify("patient_ehr")
    print(level)   # DataClass.CRITICAL
"""

from __future__ import annotations

from enum import Enum
from typing import Optional

from .policy import Policy


class DataClass(str, Enum):
    """
    Data classification levels, ordered from least to most sensitive.
    Aligned with ISO 27001 / NIST SP 800-60 classification schemes.
    """
    PUBLIC       = "PUBLIC"       # No sensitivity — public web content, docs
    INTERNAL     = "INTERNAL"     # Internal use — company wikis, internal FAQs
    CONFIDENTIAL = "CONFIDENTIAL" # Business sensitive — customer lists, contracts
    RESTRICTED   = "RESTRICTED"   # PII, financial records, employee data
    CRITICAL     = "CRITICAL"     # PHI, payment card data, biometrics, psychiatric


# ── Built-in data source classifications ──────────────────────────────────────
# Covers all data sources used in the SDK's built-in policy sets.

BUILTIN_CLASSIFICATIONS: dict[str, DataClass] = {
    # PUBLIC
    "public_sanctions_list":        DataClass.PUBLIC,
    "fatf_typologies":              DataClass.PUBLIC,
    "public_medical_literature":    DataClass.PUBLIC,
    "clinical_guidelines_db":       DataClass.PUBLIC,
    "regulatory_guidelines":        DataClass.PUBLIC,
    "return_policy_db":             DataClass.PUBLIC,
    "internal_faq":                 DataClass.INTERNAL,
    # INTERNAL
    "shipping_tracker":             DataClass.INTERNAL,
    "order_management_db":          DataClass.INTERNAL,
    "aml_alerts":                   DataClass.INTERNAL,
    "fraud_patterns":               DataClass.INTERNAL,
    "investigation_notes":          DataClass.INTERNAL,
    # CONFIDENTIAL
    "compliance_reports":           DataClass.CONFIDENTIAL,
    "support_ticket_db":            DataClass.CONFIDENTIAL,
    "loyalty_program_db":           DataClass.CONFIDENTIAL,
    "purchase_history":             DataClass.CONFIDENTIAL,
    "crm_customer_profile":         DataClass.CONFIDENTIAL,
    "drug_interaction_db":          DataClass.CONFIDENTIAL,
    "pharmacy_inventory":           DataClass.CONFIDENTIAL,
    # RESTRICTED
    "transaction_history":          DataClass.RESTRICTED,
    "banking_records":              DataClass.RESTRICTED,
    "account_balances":             DataClass.RESTRICTED,
    "wire_transfer_details":        DataClass.RESTRICTED,
    "payment_records":              DataClass.RESTRICTED,
    "kyc_documents":                DataClass.RESTRICTED,
    "customer_full_profile":        DataClass.RESTRICTED,
    "patient_vitals":               DataClass.RESTRICTED,
    "lab_results":                  DataClass.RESTRICTED,
    "imaging_results":              DataClass.RESTRICTED,
    # CRITICAL
    "patient_ehr":                  DataClass.CRITICAL,
    "clinical_history":             DataClass.CRITICAL,
    "psychiatric_records":          DataClass.CRITICAL,
}


class DataClassifier:
    """
    Classifies data sources by sensitivity level.

    Usage::

        classifier = DataClassifier()
        level = classifier.classify("patient_ehr")
        # DataClass.CRITICAL

        classifier.register("my_db", DataClass.RESTRICTED)
        level = classifier.classify("my_db")
        # DataClass.RESTRICTED

        # Highest classification among a list of sources
        max_level = classifier.max_classification(
            ["public_sanctions_list", "patient_ehr", "transaction_history"]
        )
        # DataClass.CRITICAL
    """

    # Ordered for comparison (higher index = more sensitive)
    _ORDER = [
        DataClass.PUBLIC, DataClass.INTERNAL, DataClass.CONFIDENTIAL,
        DataClass.RESTRICTED, DataClass.CRITICAL,
    ]

    def __init__(self, override: Optional[dict[str, DataClass]] = None):
        self._map: dict[str, DataClass] = {**BUILTIN_CLASSIFICATIONS}
        if override:
            self._map.update(override)

    def register(self, source_name: str, classification: DataClass) -> "DataClassifier":
        """Register or update the classification of a data source."""
        self._map[source_name.lower().replace(" ", "_")] = classification
        return self

    def classify(
        self,
        source: str,
        default: DataClass = DataClass.INTERNAL,
    ) -> DataClass:
        """Return the classification of a data source."""
        key = source.lower().replace(" ", "_").replace("-", "_")
        if key in self._map:
            return self._map[key]
        # Fuzzy match: check if any registered key is a substring
        for k, v in self._map.items():
            if k in key or key in k:
                return v
        return default

    def max_classification(
        self,
        sources: list[str],
        default: DataClass = DataClass.INTERNAL,
    ) -> DataClass:
        """Return the highest (most sensitive) classification among a list of sources."""
        if not sources:
            return default
        levels = [self.classify(s, default) for s in sources]
        return max(levels, key=lambda l: self._ORDER.index(l))

    def all_above(
        self,
        sources: list[str],
        threshold: DataClass,
    ) -> list[str]:
        """Return all sources at or above the given classification threshold."""
        idx = self._ORDER.index(threshold)
        return [s for s in sources
                if self._ORDER.index(self.classify(s)) >= idx]

    def list_all(self) -> dict[str, str]:
        """Return all registered classifications."""
        return {k: v.value for k, v in sorted(self._map.items())}


# ── Auto-generate policies from classification rules ──────────────────────────

def build_classification_policies(
    classifier:                  Optional[DataClassifier] = None,
    restricted_requires_approval: bool = True,
    critical_always_blocks:       bool = True,
    confidential_min_context:     int  = 0,
    policy_id_prefix:             str  = "DC",
) -> list[Policy]:
    """
    Generate governance policies automatically from data classification rules.

    Args:
        classifier:                   DataClassifier instance. Default = built-in.
        restricted_requires_approval: If True, accessing RESTRICTED sources requires
                                       human approval (ESCALATE).
        critical_always_blocks:       If True, accessing CRITICAL sources is always
                                       blocked (hard BLOCK).
        confidential_min_context:     Minimum context sources allowed without
                                       additional scrutiny (0 = no limit).
        policy_id_prefix:             Prefix for generated policy IDs.

    Returns:
        List of Policy objects that can be passed to ControlPlane(policies=...).
    """
    _clf = classifier or DataClassifier()
    policies: list[Policy] = []
    n = 1

    if critical_always_blocks:
        critical_sources = {k for k, v in BUILTIN_CLASSIFICATIONS.items()
                            if v == DataClass.CRITICAL}
        policies.append(Policy(
            policy_id   = f"{policy_id_prefix}-{n:03d}",
            name        = "Critical Data Source Access — Hard Block",
            description = (
                "Accessing CRITICAL data sources (PHI, psychiatric records, "
                "biometrics) is always blocked for autonomous AI agents. "
                "Requires human authorization workflow."
            ),
            category    = "DATA_GOVERNANCE",
            blocks_action           = True,
            requires_human_approval = False,
            evaluate    = lambda action, _srcs=critical_sources, _clf=_clf: (
                not bool(_clf.all_above(action.context_sources, DataClass.CRITICAL)),
                f"Access to CRITICAL data sources blocked: "
                f"{_clf.all_above(action.context_sources, DataClass.CRITICAL)}"
            ),
        ))
        n += 1

    if restricted_requires_approval:
        policies.append(Policy(
            policy_id   = f"{policy_id_prefix}-{n:03d}",
            name        = "Restricted Data Source Access — Requires Approval",
            description = (
                "Accessing RESTRICTED data sources (PII, financial records, "
                "payment data) requires human approval before AI execution."
            ),
            category    = "DATA_GOVERNANCE",
            blocks_action           = False,
            requires_human_approval = True,
            evaluate    = lambda action, _clf=_clf: (
                not bool(_clf.all_above(action.context_sources, DataClass.RESTRICTED)),
                f"Human approval required: RESTRICTED sources accessed: "
                f"{_clf.all_above(action.context_sources, DataClass.RESTRICTED)}"
            ),
        ))
        n += 1

    return policies


# ── Context classification helper ──────────────────────────────────────────────

def classify_context_sources(
    sources:    list[str],
    classifier: Optional[DataClassifier] = None,
) -> dict[str, str]:
    """
    Classify a list of context sources and return a mapping of source → class.

    Useful for auditing what data an agent is accessing.

    Example::

        classes = classify_context_sources(["patient_ehr", "public_sanctions_list"])
        # {"patient_ehr": "CRITICAL", "public_sanctions_list": "PUBLIC"}
    """
    clf = classifier or DataClassifier()
    return {src: clf.classify(src).value for src in sources}
