"""
encyclopedia_rules.py
─────────────────────
Integration knowledge derived from the System Design Components Encyclopedia.
Used to enrich the LLM system prompt and extend the family → concrete mapping
to cover new tasks (data_pipeline, security_hardened_api, observability_stack).

Import this in round2_inf.py and inference.py to replace / extend:
  - SYSTEM_PROMPT
  - FAMILY_MAP / TASK_BROKER_MAP / TASK_BONUS_TARGETS
"""

from __future__ import annotations
from typing import Dict, List

# ---------------------------------------------------------------------------
# Extended family -> concrete component map
# New families from encyclopedia: api_gateway, waf, iam, secrets_manager,
# logging, tracing, alerting, stream_processor, etl_orchestrator,
# graph_database, time_series_db, data_warehouse, serverless, circuit_breaker
# ---------------------------------------------------------------------------

FAMILY_MAP: Dict[str, str] = {
    # original families
    "database":              "postgres",
    "cache":                 "redis",
    "broker":                "rabbitmq",      # overridden per task
    "worker":                "worker",
    "websocket_gateway":     "websocket_gateway",
    "search":                "elasticsearch",
    "storage":               "s3",
    "cdn":                   "cloudfront",
    "transcoder_worker":     "transcoder_worker",
    "recommendation_worker": "recommendation_worker",
    "load_balancer":         "nginx",
    "auth":                  "keycloak",
    "observability":         "prometheus",
    "geospatial_index":      "geospatial_index",
    "feature_store":         "feast",
    "model_registry":        "mlflow",
    "inference_server":      "triton",
    "presence_service":      "presence_service",
    "notification_service":  "notification_service",
    "payment_gateway":       "stripe",
    "rate_limiting":         "kong",
    "metadata_service":      "metadata_service",
    # new families from encyclopedia
    "api_gateway":           "kong",
    "waf":                   "aws_waf",
    "iam":                   "aws_iam",
    "secrets_manager":       "vault",
    "logging":               "fluentd",
    "tracing":               "jaeger",
    "alerting":              "pagerduty",
    "stream_processor":      "flink",
    "etl_orchestrator":      "airflow",
    "data_warehouse":        "redshift",
    "graph_database":        "neo4j",
    "time_series_db":        "influxdb",
    "serverless":            "lambda",
    "circuit_breaker":       "resilience4j",
}

TASK_BROKER_MAP: Dict[str, str] = {
    "youtube_platform":      "kafka",
    "ride_sharing":          "kafka",
    "ml_platform":           "kafka",
    "data_pipeline":         "kafka",
    "observability_stack":   "kafka",
    "chat_system":           "rabbitmq",
    "ecommerce_platform":    "rabbitmq",
    "url_shortener":         "rabbitmq",
    "security_hardened_api": "rabbitmq",
}

TASK_BONUS_TARGETS: Dict[str, List[str]] = {
    "chat_system":            ["keycloak", "prometheus", "presence_service", "notification_service"],
    "ecommerce_platform":     ["keycloak", "prometheus", "kong"],
    "youtube_platform":       ["keycloak", "prometheus", "recommendation_worker"],
    "ride_sharing":           ["keycloak", "prometheus", "stripe", "notification_service"],
    "url_shortener":          ["keycloak", "prometheus", "kong"],
    "ml_platform":            ["keycloak", "prometheus"],
    "data_pipeline":          ["jaeger", "pagerduty", "fluentd"],
    "security_hardened_api":  ["prometheus", "splunk", "pagerduty", "jaeger"],
    "observability_stack":    ["influxdb", "elasticsearch"],
}


# ---------------------------------------------------------------------------
# Enriched system prompt — injects encyclopedia integration rules
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = """You are a senior software architect in an incremental
architecture-design game. Each turn output EXACTLY ONE action — nothing else.

Valid actions
─────────────
  add <component>           e.g.  add postgres
  connect <comp1> <comp2>   e.g.  connect kafka flink
  submit                    finalise the design

Core rules
──────────
1. ONE action per reply, lowercase, no punctuation, no explanation.
2. Add a component before connecting it.
3. NEVER repeat an action already done. If rejected, switch strategy.
4. Finish ALL required items and connections before adding bonus items.
5. Submit only when missing_required_items=[] AND missing_required_connections=[].

Family → concrete component (task-aware)
─────────────────────────────────────────
  database          → postgres
  cache             → redis
  search            → elasticsearch
  storage           → s3
  cdn               → cloudfront
  load_balancer     → nginx
  auth              → keycloak
  observability     → prometheus
  rate_limiting     → kong
  payment_gateway   → stripe
  geospatial_index  → geospatial_index
  feature_store     → feast
  model_registry    → mlflow
  inference_server  → triton
  worker            → worker
  websocket_gateway → websocket_gateway

  NEW (from encyclopedia):
  api_gateway       → kong
  waf               → aws_waf
  iam               → aws_iam
  secrets_manager   → vault
  logging           → fluentd
  tracing           → jaeger
  alerting          → pagerduty
  stream_processor  → flink
  etl_orchestrator  → airflow
  data_warehouse    → redshift
  graph_database    → neo4j
  time_series_db    → influxdb
  serverless        → lambda
  circuit_breaker   → resilience4j

  Broker is task-dependent:
    chat_system / ecommerce_platform / url_shortener / security_hardened_api → rabbitmq
    youtube_platform / ride_sharing / ml_platform / data_pipeline / observability_stack → kafka

Encyclopedia integration patterns (connect these pairs when required)
──────────────────────────────────────────────────────────────────────
  Edge layer      : waf → api_gateway → load_balancer → api_server
  Auth flow       : api_gateway → auth → database
  Secrets         : api_server → secrets_manager (fetch creds before DB connect)
  Cache pattern   : api_server → cache → database
  Async pattern   : api_server → broker → worker → database
  Streaming       : api_server → broker → stream_processor → data_warehouse
  Observability   : api_server → observability; observability → alerting
  Log pipeline    : broker → logging → storage (cold archive)
  Trace pipeline  : broker → tracing → observability
  ML pipeline     : broker → feature_store → inference_server; model_registry → inference_server
  Video pipeline  : api_server → broker → transcoder_worker → storage; api_server → cdn

After all required work is complete, add bonus items then submit.
"""


def effective_map(task_name: str) -> Dict[str, str]:
    """Return family map with task-aware broker override."""
    m = dict(FAMILY_MAP)
    m["broker"] = TASK_BROKER_MAP.get(task_name, "rabbitmq")
    return m