#!/usr/bin/env python3
"""Seed model pricing for Cohorta (Azure AI Foundry).

Run from the backend/ directory:
    source venv/bin/activate
    python scripts/seed_model_pricing.py

Idempotent — upserts on {model, effective_from}. Safe to re-run.
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from dotenv import load_dotenv
load_dotenv()

import pymongo
from datetime import datetime, timezone

MONGO_URI = os.environ.get("MONGO_URI")
MONGO_USER = os.environ.get("MONGO_USER")
MONGO_PASS = os.environ.get("MONGO_PASS")
MONGO_HOST = os.environ.get("MONGO_HOST", "localhost")
MONGO_PORT = os.environ.get("MONGO_PORT", "27017")
MONGO_DB = os.environ.get("MONGO_DB", "cohorta_db")

if not MONGO_URI:
    if MONGO_USER and MONGO_PASS:
        MONGO_URI = f"mongodb://{MONGO_USER}:{MONGO_PASS}@{MONGO_HOST}:{MONGO_PORT}/{MONGO_DB}?authSource=admin"
    else:
        MONGO_URI = f"mongodb://{MONGO_HOST}:{MONGO_PORT}"

# Pricing effective from project launch
EFFECTIVE_FROM = datetime(2026, 5, 23, tzinfo=timezone.utc)

# Azure AI Foundry pricing (GlobalStandard, <272k context, USD per 1M tokens)
# Source: Azure OpenAI pricing page, May 2026
PRICING_ROWS = [
    {
        "model": "gpt-5.4",
        "provider": "azure",
        "currency": "USD",
        "tiers": [
            {
                "threshold_input_tokens": 0,
                "input_per_mtok": 2.50,
                "cached_input_per_mtok": 0.25,
                "output_per_mtok": 15.00,
                "image_per_mtok": None,
            }
        ],
        "effective_from": EFFECTIVE_FROM,
        "effective_until": None,
        "notes": "gpt-5.4 GlobalStandard <272k ctx. Retirement 2027-03-05.",
    },
    {
        "model": "gpt-5.4-mini",
        "provider": "azure",
        "currency": "USD",
        "tiers": [
            {
                "threshold_input_tokens": 0,
                "input_per_mtok": 0.75,
                "cached_input_per_mtok": 0.08,
                "output_per_mtok": 4.50,
                "image_per_mtok": None,
            }
        ],
        "effective_from": EFFECTIVE_FROM,
        "effective_until": None,
        "notes": "gpt-5.4-mini GlobalStandard. Used for cheap features (summary, key_themes, etc.). Retirement 2027-03-18.",
    },
    {
        "model": "gpt-5.4-nano",
        "provider": "azure",
        "currency": "USD",
        "tiers": [
            {
                "threshold_input_tokens": 0,
                "input_per_mtok": 0.20,
                "cached_input_per_mtok": 0.02,
                "output_per_mtok": 1.25,
                "image_per_mtok": None,
            }
        ],
        "effective_from": EFFECTIVE_FROM,
        "effective_until": None,
        "notes": "gpt-5.4-nano GlobalStandard. Optional ultra-cheap tier.",
    },
]


def main():
    client = pymongo.MongoClient(MONGO_URI, serverSelectionTimeoutMS=5000)
    db = client[MONGO_DB]

    db.model_pricing.create_index(
        [("model", pymongo.ASCENDING), ("effective_from", pymongo.DESCENDING)],
        background=True,
    )

    for row in PRICING_ROWS:
        key = {"model": row["model"], "effective_from": row["effective_from"]}
        result = db.model_pricing.update_one(key, {"$set": row}, upsert=True)
        action = "inserted" if result.upserted_id else "updated"
        print(f"  {action}: {row['model']} (effective from {row['effective_from'].date()})")

    print(f"\nDone. {len(PRICING_ROWS)} pricing rows seeded.")
    client.close()


if __name__ == "__main__":
    main()