"""AISA layer and cross-layer contract definitions.

Defines the 7 architectural layers and 4 cross-layer contracts
from the AISA reference architecture, along with evaluation criteria.
"""

from dataclasses import dataclass, field


@dataclass
class Criterion:
    """A single evaluation criterion within a layer."""
    id: str
    name: str
    description: str
    weight: float = 1.0


@dataclass
class Layer:
    """An AISA architectural layer."""
    id: str
    name: str
    short_name: str
    description: str
    criteria: list[Criterion] = field(default_factory=list)


@dataclass
class Contract:
    """An AISA cross-layer contract."""
    id: str
    name: str
    description: str
    criteria: list[Criterion] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Layer Definitions
# ---------------------------------------------------------------------------

L1_LLM_FOUNDATION = Layer(
    id="L1",
    name="LLM Foundation Layer",
    short_name="LLM Foundation",
    description="Core language understanding, reasoning, and inference substrate.",
    criteria=[
        Criterion("L1.1", "Model Adapters", "Supports multiple LLM providers or model swapping"),
        Criterion("L1.2", "Prompt Engineering", "Structured prompt templates with variable interpolation"),
        Criterion("L1.3", "Context Window Management", "Token budgeting, truncation, or summarization strategies"),
        Criterion("L1.4", "Safety Filters & Grounding", "Content filtering, refusal policies, or grounding mechanisms"),
        Criterion("L1.5", "Instruction Tuning / Alignment", "Fine-tuning, RLHF, or alignment configuration"),
        Criterion("L1.6", "Output Structure", "Structured outputs (JSON, Pydantic, typed responses)"),
    ],
)

L2_TOOL_ENVIRONMENT = Layer(
    id="L2",
    name="Tool & Environment Layer",
    short_name="Tool & Environment",
    description="Agent's actuation interface for executing external actions safely.",
    criteria=[
        Criterion("L2.1", "Tool Definitions", "Formal tool/function definitions with schemas"),
        Criterion("L2.2", "Schema Validation", "Input/output validation on tool calls"),
        Criterion("L2.3", "Sandboxing", "Isolated execution environment for code/tools"),
        Criterion("L2.4", "Permission Controls", "Per-tool or per-agent permission scoping"),
        Criterion("L2.5", "MCP Support", "Model Context Protocol integration"),
        Criterion("L2.6", "Error Handling & Retries", "Structured error propagation and retry logic"),
        Criterion("L2.7", "Rate Limiting", "API rate limiting or throttling controls"),
    ],
)

L3_COGNITIVE_AGENT = Layer(
    id="L3",
    name="Cognitive Agent Layer",
    short_name="Cognitive Agent",
    description="Agent reasoning, planning, reflection, and self-improvement.",
    criteria=[
        Criterion("L3.1", "Task Planning & Decomposition", "Explicit planning or task breakdown capabilities"),
        Criterion("L3.2", "Multi-turn Reasoning", "Coherent reasoning across multiple interaction turns"),
        Criterion("L3.3", "Memory Management", "Short-term, long-term, or episodic memory systems"),
        Criterion("L3.4", "Goal Tracking", "Explicit goal state tracking and progress monitoring"),
        Criterion("L3.5", "Reflection Loops", "Self-critique, self-evaluation, or improvement loops"),
        Criterion("L3.6", "External Feedback Integration", "Incorporating human or system feedback into reasoning"),
    ],
)

L4_INFRASTRUCTURE = Layer(
    id="L4",
    name="Agentic Infrastructure Layer",
    short_name="Infrastructure",
    description="Coordination, orchestration, and scalability backbone.",
    criteria=[
        Criterion("L4.1", "Workflow Orchestration", "DAG, state machine, or pipeline orchestration"),
        Criterion("L4.2", "State Management", "Persistent or checkpointable agent state"),
        Criterion("L4.3", "Multi-Agent Coordination", "Agent-to-agent communication and delegation"),
        Criterion("L4.4", "Observability & Tracing", "Distributed tracing, span propagation, dashboards"),
        Criterion("L4.5", "Logging & Monitoring", "Structured logging with correlation IDs"),
        Criterion("L4.6", "Agent Lifecycle Management", "Agent creation, scaling, and teardown"),
        Criterion("L4.7", "Latency & Cost Optimization", "Caching, batching, routing optimizations"),
    ],
)

L5_EVALUATION = Layer(
    id="L5",
    name="Evaluation & Feedback Layer",
    short_name="Evaluation",
    description="Continuous behavioral assurance and validation framework.",
    criteria=[
        Criterion("L5.1", "Component-Level Evaluation", "Unit tests for prompts, tools, retrieval components"),
        Criterion("L5.2", "Trajectory-Level Evaluation", "End-to-end evaluation of multi-step agent behavior"),
        Criterion("L5.3", "Behavioral Monitoring", "Runtime drift detection and anomaly alerting"),
        Criterion("L5.4", "Regression Testing", "Automated tests that catch behavioral regressions"),
        Criterion("L5.5", "Human-in-the-Loop Evaluation", "Structured human review workflows"),
        Criterion("L5.6", "Quality Metrics", "Accuracy, reliability, interpretability metrics"),
    ],
)

L6_DEPLOYMENT = Layer(
    id="L6",
    name="Development & Deployment Layer",
    short_name="Dev & Deployment",
    description="Operational lifecycle for agentic systems as production software.",
    criteria=[
        Criterion("L6.1", "Version Control of Agents", "Versioned agent configs, prompts, tools, workflows"),
        Criterion("L6.2", "CI/CD Integration", "Automated build, test, deploy pipelines"),
        Criterion("L6.3", "A/B Testing", "Traffic splitting and experiment management"),
        Criterion("L6.4", "Performance Benchmarking", "Latency, cost, quality baseline tracking"),
        Criterion("L6.5", "Staged Rollout", "Canary, shadow, or blue-green deployment strategies"),
        Criterion("L6.6", "Secrets Management", "Secure handling of API keys and credentials"),
        Criterion("L6.7", "Model Lifecycle Tracking", "Tracking model versions, adapters, and dependencies"),
    ],
)

L7_GOVERNANCE = Layer(
    id="L7",
    name="Governance, Ethics & Policy Layer",
    short_name="Governance",
    description="System-wide constraints aligned with human values and regulations.",
    criteria=[
        Criterion("L7.1", "Policy-as-Code", "Machine-enforceable governance policies"),
        Criterion("L7.2", "Privacy & Data Protection", "PII detection, redaction, data retention policies"),
        Criterion("L7.3", "Fairness & Bias Mitigation", "Bias detection, fairness auditing"),
        Criterion("L7.4", "Accountability Framework", "Audit trails, ownership assignment, escalation paths"),
        Criterion("L7.5", "Human Oversight", "Approval gates, intervention mechanisms, escalation"),
        Criterion("L7.6", "Regulatory Compliance", "Alignment with GDPR, HIPAA, SOC2, NIST AI RMF"),
        Criterion("L7.7", "Explainability", "Decision transparency and reasoning traces"),
    ],
)

# ---------------------------------------------------------------------------
# Cross-Layer Contracts
# ---------------------------------------------------------------------------

C1_POLICY = Contract(
    id="C1",
    name="Policy Contract",
    description="Enforces governance and safety constraints across all layers.",
    criteria=[
        Criterion("C1.1", "Cross-Layer Policy Propagation", "Policies cascade across tool, cognitive, and infra layers"),
        Criterion("C1.2", "Tool Allowlists/Denylists", "Enforceable tool permission scopes"),
        Criterion("C1.3", "Data Classification Rules", "Data sensitivity labels and handling rules"),
    ],
)

C2_TELEMETRY = Contract(
    id="C2",
    name="Telemetry Contract",
    description="Enables auditability and reproducibility across layers.",
    criteria=[
        Criterion("C2.1", "End-to-End Trace IDs", "Correlation IDs spanning all layers"),
        Criterion("C2.2", "Structured Action Logs", "Tool calls, decisions, and outcomes logged consistently"),
        Criterion("C2.3", "OpenTelemetry Compatibility", "OTel-compatible telemetry emission"),
    ],
)

C3_VERSIONING = Contract(
    id="C3",
    name="Versioning Contract",
    description="Ensures reproducibility and controlled evolution.",
    criteria=[
        Criterion("C3.1", "Agent Package Versioning", "Agent as a versioned artifact (model + prompts + tools)"),
        Criterion("C3.2", "Prompt & Tool Versioning", "Individual prompt and tool schema version tracking"),
        Criterion("C3.3", "Reproducible Runs", "Any output traceable to specific artifact versions"),
    ],
)

C4_BUDGET = Contract(
    id="C4",
    name="Budget Contract",
    description="Controls cost, latency, and resource consumption.",
    criteria=[
        Criterion("C4.1", "Token Budget Enforcement", "Hard limits on token consumption per run/agent"),
        Criterion("C4.2", "Cost Ceiling Controls", "Monetary cost caps with circuit-breaking"),
        Criterion("C4.3", "Action Budget Limits", "Limits on tool calls, retries, or steps per task"),
    ],
)

# ---------------------------------------------------------------------------
# Registries
# ---------------------------------------------------------------------------

ALL_LAYERS = [
    L1_LLM_FOUNDATION,
    L2_TOOL_ENVIRONMENT,
    L3_COGNITIVE_AGENT,
    L4_INFRASTRUCTURE,
    L5_EVALUATION,
    L6_DEPLOYMENT,
    L7_GOVERNANCE,
]

ALL_CONTRACTS = [
    C1_POLICY,
    C2_TELEMETRY,
    C3_VERSIONING,
    C4_BUDGET,
]


def get_total_criteria_count() -> int:
    """Return total number of evaluation criteria across all layers and contracts."""
    total = sum(len(layer.criteria) for layer in ALL_LAYERS)
    total += sum(len(contract.criteria) for contract in ALL_CONTRACTS)
    return total
