"""AISA gap-specific recommendations.

Provides actionable recommendations for each AISA layer and contract
based on detected gaps and the frameworks in use.
"""

LAYER_RECOMMENDATIONS: dict[str, dict] = {
    "L1": {
        "name": "LLM Foundation Layer",
        "priority": "Medium",
        "gaps": {
            "L1.1": "Add model adapter abstraction (e.g., LiteLLM) to avoid vendor lock-in and enable fallback strategies.",
            "L1.2": "Implement structured prompt templates with version control. Consider Jinja2 or framework-native templating.",
            "L1.3": "Add token counting and context window management. Implement truncation or summarization strategies for long conversations.",
            "L1.4": "Add safety filters (content moderation API, guardrails library) and grounding via RAG or citation requirements.",
            "L1.5": "Consider domain-specific fine-tuning or alignment if generic models underperform on your use case.",
            "L1.6": "Use structured output formats (Pydantic models, JSON schema) to improve reliability and downstream processing.",
        },
    },
    "L2": {
        "name": "Tool & Environment Layer",
        "priority": "High",
        "gaps": {
            "L2.1": "Define all tools with formal schemas (name, description, parameters, return types). Use framework-native tool decorators.",
            "L2.2": "Add Pydantic or JSON Schema validation on tool inputs and outputs to catch malformed calls.",
            "L2.3": "Implement sandboxing for code execution (Docker, gVisor, or restricted Python). Never run untrusted code in-process.",
            "L2.4": "Implement per-agent permission scopes. Use allowlists to restrict which tools each agent can access.",
            "L2.5": "Consider adopting MCP (Model Context Protocol) for standardized tool discovery and invocation.",
            "L2.6": "Add structured error handling with retry logic, backoff, and circuit-breaking for tool failures.",
            "L2.7": "Implement rate limiting on external API calls to prevent quota exhaustion and runaway costs.",
        },
    },
    "L3": {
        "name": "Cognitive Agent Layer",
        "priority": "Medium",
        "gaps": {
            "L3.1": "Add explicit planning capabilities (Plan-and-Execute, hierarchical task decomposition).",
            "L3.2": "Ensure chat history is maintained and managed across turns. Implement summarization for long sessions.",
            "L3.3": "Implement layered memory: short-term (buffer), long-term (vector store), and episodic (session logs).",
            "L3.4": "Add explicit goal tracking with progress indicators and termination criteria.",
            "L3.5": "Implement reflection loops where the agent self-evaluates outputs before returning them.",
            "L3.6": "Add human-in-the-loop hooks for feedback during execution, not just at the end.",
        },
    },
    "L4": {
        "name": "Agentic Infrastructure Layer",
        "priority": "High",
        "gaps": {
            "L4.1": "Implement workflow orchestration (LangGraph, Temporal, or custom state machines) for predictable execution.",
            "L4.2": "Add state checkpointing to enable recovery from failures and support long-running tasks.",
            "L4.3": "Design multi-agent coordination patterns (supervisor, delegation, handoff) if your system uses multiple agents.",
            "L4.4": "Add distributed tracing (OpenTelemetry) to track requests across all components.",
            "L4.5": "Implement structured logging with correlation IDs that span the full request lifecycle.",
            "L4.6": "Implement agent lifecycle management for creation, health checks, scaling, and graceful shutdown.",
            "L4.7": "Add caching for repeated LLM calls, batching for parallel operations, and smart routing for cost optimization.",
        },
    },
    "L5": {
        "name": "Evaluation & Feedback Layer",
        "priority": "Critical",
        "gaps": {
            "L5.1": "Create unit tests for individual prompts, tools, and retrieval components.",
            "L5.2": "Build end-to-end evaluation suites that test complete agent trajectories against expected outcomes.",
            "L5.3": "Implement behavioral monitoring to detect drift, anomalies, or degradation in production.",
            "L5.4": "Add regression tests with golden datasets to catch behavioral regressions from prompt/model changes.",
            "L5.5": "Set up structured human evaluation workflows with calibrated rubrics and inter-rater agreement tracking.",
            "L5.6": "Define and track quality metrics (accuracy, faithfulness, relevancy, latency, cost per task).",
        },
    },
    "L6": {
        "name": "Development & Deployment Layer",
        "priority": "Critical",
        "gaps": {
            "L6.1": "Version all agent artifacts (prompts, tool schemas, workflows, configs) in source control.",
            "L6.2": "Set up CI/CD pipelines that run evaluation suites before deploying agent changes.",
            "L6.3": "Implement A/B testing to compare agent variants in production with statistical rigor.",
            "L6.4": "Establish performance baselines (latency p50/p95, cost per task, quality scores) and track them over time.",
            "L6.5": "Use staged rollout (canary/shadow) to limit blast radius when deploying agent changes.",
            "L6.6": "Use a secrets manager (HashiCorp Vault, AWS Secrets Manager, or env vars with .env) for all credentials.",
            "L6.7": "Track which model versions, adapters, and dependencies are used in each deployment.",
        },
    },
    "L7": {
        "name": "Governance, Ethics & Policy Layer",
        "priority": "Critical",
        "gaps": {
            "L7.1": "Implement policy-as-code: define governance rules in machine-readable format and enforce them at runtime.",
            "L7.2": "Add PII detection and redaction. Implement data retention policies and deletion workflows.",
            "L7.3": "Audit agent outputs for fairness and bias. Consider using fairness toolkits (Fairlearn, AIF360).",
            "L7.4": "Create audit trails with ownership assignment, escalation paths, and incident response procedures.",
            "L7.5": "Implement approval gates for high-risk actions. Define escalation thresholds and human review triggers.",
            "L7.6": "Map your system to relevant regulatory frameworks (NIST AI RMF, EU AI Act, GDPR, HIPAA as applicable).",
            "L7.7": "Add reasoning traces and decision logs that can be reviewed by stakeholders and auditors.",
        },
    },
}

CONTRACT_RECOMMENDATIONS: dict[str, dict] = {
    "C1": {
        "name": "Policy Contract",
        "priority": "Critical",
        "gaps": {
            "C1.1": "Design a policy engine that propagates constraints across tool, cognitive, and infrastructure layers.",
            "C1.2": "Implement tool allowlists/denylists per agent, per environment, and per risk level.",
            "C1.3": "Classify data by sensitivity level (public, internal, confidential, restricted) with automated handling rules.",
        },
    },
    "C2": {
        "name": "Telemetry Contract",
        "priority": "High",
        "gaps": {
            "C2.1": "Generate unique trace/correlation IDs at request entry and propagate them through all layers.",
            "C2.2": "Log all tool calls, agent decisions, and outcomes in a structured, queryable format.",
            "C2.3": "Emit telemetry in OpenTelemetry format to enable integration with standard observability platforms.",
        },
    },
    "C3": {
        "name": "Versioning Contract",
        "priority": "High",
        "gaps": {
            "C3.1": "Package agents as versioned artifacts that bundle model, prompts, tools, and configs together.",
            "C3.2": "Version prompts and tool schemas independently so changes can be tracked and compared.",
            "C3.3": "Ensure any past agent run can be reproduced from its version snapshot for debugging and compliance.",
        },
    },
    "C4": {
        "name": "Budget Contract",
        "priority": "Critical",
        "gaps": {
            "C4.1": "Enforce token budgets per agent run with hard stops when limits are exceeded.",
            "C4.2": "Set monetary cost ceilings with circuit-breaking that gracefully terminates over-budget runs.",
            "C4.3": "Limit the number of tool calls, retries, and reasoning steps to prevent infinite loops.",
        },
    },
}


def get_recommendation(criterion_id: str) -> str | None:
    """Get the recommendation for a specific gap."""
    layer_id = criterion_id.split(".")[0]

    if layer_id.startswith("L"):
        layer_recs = LAYER_RECOMMENDATIONS.get(layer_id, {})
        return layer_recs.get("gaps", {}).get(criterion_id)
    elif layer_id.startswith("C"):
        contract_recs = CONTRACT_RECOMMENDATIONS.get(layer_id, {})
        return contract_recs.get("gaps", {}).get(criterion_id)

    return None


def get_priority(component_id: str) -> str:
    """Get the priority level for a layer or contract."""
    if component_id in LAYER_RECOMMENDATIONS:
        return LAYER_RECOMMENDATIONS[component_id]["priority"]
    if component_id in CONTRACT_RECOMMENDATIONS:
        return CONTRACT_RECOMMENDATIONS[component_id]["priority"]
    return "Medium"
