# ---------------------------------------------------------------------------
#  Text-to-Cosmos (NoSQL) — minimal working skeleton  (Python)
#  Console script. Azure OpenAI generates the query, Cosmos SDK runs it.
#
#  Install:
#     pip install openai azure-cosmos
#
#  Teaching skeleton: one file, minimal error handling, hard-coded sample
#  schema. Get it running, then refactor.
# ---------------------------------------------------------------------------

from openai import AzureOpenAI
from azure.cosmos import CosmosClient

# ── 1. CONFIG ──────────────────────────────────────────────────────────────
# In real life read these from environment variables, never hard-code them.
# The Cosmos key here should be a READ-ONLY key.
AOAI_ENDPOINT   = "https://YOUR-RESOURCE.openai.azure.com/"
AOAI_KEY        = "YOUR_AZURE_OPENAI_KEY"
AOAI_DEPLOYMENT = "gpt-4o"          # your *deployment* name, not the model name
AOAI_API_VER    = "2024-10-21"      # a current Azure OpenAI API version

COSMOS_ENDPOINT = "https://YOUR-COSMOS.documents.azure.com:443/"
COSMOS_READ_KEY = "YOUR_COSMOS_READONLY_KEY"
DATABASE_NAME   = "SalesDb"
CONTAINER_NAME  = "Orders"

# ── 2. THE PROMPT — this is where 80% of the quality lives ──────────────────
# Three ingredients: (a) your document shape, (b) the Cosmos dialect rules,
# (c) a few worked examples. Change these to match YOUR data and the whole
# engine adapts — no code changes needed.
SYSTEM_PROMPT = """
You translate a user's plain-English question into a single Azure Cosmos DB
NoSQL (SQL API) query. Return ONLY the query text — no explanation, no markdown.

THE DOCUMENTS look like this (container alias is always 'c'):
{
  "id": "string",
  "customerName": "string",
  "total": 123.45,
  "orderDate": "2026-08-14",     // ISO date string
  "status": "shipped",           // pending | shipped | cancelled
  "items": [ { "sku": "string", "qty": 2, "price": 9.99 } ]
}

DIALECT RULES (Cosmos NoSQL is SQL-LIKE, not T-SQL — obey these):
- SELECT queries only. Never INSERT/UPDATE/DELETE.
- Reference the container as 'c'. e.g. SELECT * FROM c WHERE c.total > 100
- No TOP, no GETDATE(). Use OFFSET/LIMIT for paging (e.g. OFFSET 0 LIMIT 10).
- JOIN is only for arrays INSIDE a document (e.g. JOIN i IN c.items),
  never across documents/containers.
- Dates are ISO strings, so compare them as strings: c.orderDate >= "2026-08-01".

EXAMPLES:
Q: orders over 500 dollars
A: SELECT * FROM c WHERE c.total > 500

Q: how many cancelled orders are there
A: SELECT VALUE COUNT(1) FROM c WHERE c.status = "cancelled"

Q: the 5 biggest orders from August 2026
A: SELECT * FROM c WHERE c.orderDate >= "2026-08-01" AND c.orderDate <= "2026-08-31" ORDER BY c.total DESC OFFSET 0 LIMIT 5

Q: customers who bought sku ABC-123
A: SELECT DISTINCT c.customerName FROM c JOIN i IN c.items WHERE i.sku = "ABC-123"
"""

# ── 3. ASK THE MODEL FOR A QUERY ────────────────────────────────────────────
question = input("Ask a question: ")

aoai = AzureOpenAI(
    azure_endpoint=AOAI_ENDPOINT,
    api_key=AOAI_KEY,
    api_version=AOAI_API_VER,
)

completion = aoai.chat.completions.create(
    model=AOAI_DEPLOYMENT,
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": question},
    ],
)

query = completion.choices[0].message.content.strip()

# Models sometimes wrap output in ```sql fences despite instructions — strip them.
query = query.strip("`").replace("sql\n", "").strip()

print(f"\nGenerated query:\n  {query}\n")

# ── 4. VALIDATE (cheap but essential) ───────────────────────────────────────
# Read-only key already protects your data; this is a second belt.
if not query.lstrip().upper().startswith("SELECT"):
    print("Rejected: not a SELECT query. Aborting.")
    raise SystemExit

# ── 5. EXECUTE AGAINST COSMOS ───────────────────────────────────────────────
cosmos = CosmosClient(COSMOS_ENDPOINT, credential=COSMOS_READ_KEY)
container = (
    cosmos.get_database_client(DATABASE_NAME)
          .get_container_client(CONTAINER_NAME)
)

print("Results:")
results = list(container.query_items(query=query))
for item in results:
    print("  ", item)

# RU cost lives in the last response headers — worth watching while you experiment.
ru = container.client_connection.last_response_headers.get("x-ms-request-charge")
print(f"\n(Request charge: {ru} RUs)")

# ── 6. (OPTIONAL) phrase the answer in English ──────────────────────────────
# Feed `results` back to the model with the original question and ask it to
# summarise. Same chat.completions.create call as step 3.

# ---------------------------------------------------------------------------
#  VERSION NOTES:
#  - Azure OpenAI is accessed through the official `openai` package via the
#    AzureOpenAI class (shown here). Bump AOAI_API_VER if a feature you want
#    needs a newer API version.
#  - azure-cosmos 4.x runs cross-partition queries automatically. Older code
#    you'll see online passes enable_cross_partition_query=True to query_items;
#    recent versions dropped that argument, so it's omitted here.
# ---------------------------------------------------------------------------