import requests
import sqlite3
import sys

DB_PATH = "tarkov.db"
API_URL = "https://api.tarkov.dev/graphql"

GRAPHQL_QUERY = """
{
  ammo {
    item {
      id
      name
      shortName
      gridImageLink
      wikiLink
    }
    caliber
    damage
    armorDamage
    penetrationPower
    penetrationChance
    fragmentationChance
    ricochetChance
    initialSpeed
    lightBleedModifier
    heavyBleedModifier
    projectileCount
    tracer
    tracerColor
    ammoType
    accuracyModifier
    recoilModifier
    staminaBurnPerDamage
  }
}
"""


def fetch_ammo():
    response = requests.post(
        API_URL,
        json={"query": GRAPHQL_QUERY},
        timeout=30
    )
    response.raise_for_status()
    data = response.json()

    if "errors" in data:
        raise RuntimeError(data["errors"])

    return data["data"]["ammo"]


def ensure_schema(conn):
    conn.executescript("""
        CREATE TABLE IF NOT EXISTS ammo (
            id                      TEXT PRIMARY KEY,
            name                    TEXT NOT NULL,
            short_name              TEXT,
            caliber                 TEXT,
            damage                  INTEGER,
            armor_damage            INTEGER,
            penetration_power       INTEGER,
            penetration_chance      REAL,
            fragmentation_chance    REAL,
            ricochet_chance         REAL,
            initial_speed           REAL,
            light_bleed_modifier    REAL,
            heavy_bleed_modifier    REAL,
            projectile_count        INTEGER DEFAULT 1,
            tracer                  INTEGER DEFAULT 0,
            tracer_color            TEXT,
            ammo_type               TEXT,
            accuracy_modifier       REAL,
            recoil_modifier         REAL,
            stamina_burn_per_damage REAL,
            grid_image_url          TEXT,
            wiki_url                TEXT,
            imported_at             TEXT DEFAULT CURRENT_TIMESTAMP,
            updated_at              TEXT DEFAULT CURRENT_TIMESTAMP
        );

        CREATE INDEX IF NOT EXISTS idx_ammo_caliber    ON ammo(caliber);
        CREATE INDEX IF NOT EXISTS idx_ammo_damage     ON ammo(damage);
        CREATE INDEX IF NOT EXISTS idx_ammo_pen        ON ammo(penetration_power);
    """)
    conn.commit()


def upsert_ammo(conn, rounds):
    inserted = 0
    updated = 0
    skipped = 0

    cursor = conn.cursor()

    for r in rounds:
        item = r.get("item") or {}
        api_id = item.get("id")
        name = item.get("name")

        if not api_id or not name:
            skipped += 1
            continue

        row_data = (
            name,
            item.get("shortName"),
            r.get("caliber"),
            r.get("damage"),
            r.get("armorDamage"),
            r.get("penetrationPower"),
            r.get("penetrationChance"),
            r.get("fragmentationChance"),
            r.get("ricochetChance"),
            r.get("initialSpeed"),
            r.get("lightBleedModifier"),
            r.get("heavyBleedModifier"),
            r.get("projectileCount") or 1,
            1 if r.get("tracer") else 0,
            r.get("tracerColor"),
            r.get("ammoType"),
            r.get("accuracyModifier"),
            r.get("recoilModifier"),
            r.get("staminaBurnPerDamage"),
            item.get("gridImageLink"),
            item.get("wikiLink"),
        )

        cursor.execute("SELECT id FROM ammo WHERE id = ?", (api_id,))
        exists = cursor.fetchone()

        if exists:
            cursor.execute("""
                UPDATE ammo SET
                    name = ?, short_name = ?, caliber = ?,
                    damage = ?, armor_damage = ?, penetration_power = ?,
                    penetration_chance = ?, fragmentation_chance = ?, ricochet_chance = ?,
                    initial_speed = ?, light_bleed_modifier = ?, heavy_bleed_modifier = ?,
                    projectile_count = ?, tracer = ?, tracer_color = ?, ammo_type = ?,
                    accuracy_modifier = ?, recoil_modifier = ?, stamina_burn_per_damage = ?,
                    grid_image_url = ?, wiki_url = ?,
                    updated_at = CURRENT_TIMESTAMP
                WHERE id = ?
            """, (*row_data, api_id))
            updated += 1
        else:
            cursor.execute("""
                INSERT INTO ammo (
                    id, name, short_name, caliber,
                    damage, armor_damage, penetration_power,
                    penetration_chance, fragmentation_chance, ricochet_chance,
                    initial_speed, light_bleed_modifier, heavy_bleed_modifier,
                    projectile_count, tracer, tracer_color, ammo_type,
                    accuracy_modifier, recoil_modifier, stamina_burn_per_damage,
                    grid_image_url, wiki_url
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (api_id, *row_data))
            inserted += 1

    conn.commit()
    return inserted, updated, skipped


def main():
    print("Fetching ammo data from tarkov.dev...")
    rounds = fetch_ammo()
    print(f"  Got {len(rounds)} rounds")

    conn = sqlite3.connect(DB_PATH)
    try:
        ensure_schema(conn)
        inserted, updated, skipped = upsert_ammo(conn, rounds)
        print(f"  Inserted: {inserted}  Updated: {updated}  Skipped: {skipped}")
    finally:
        conn.close()

    print("Done.")


if __name__ == "__main__":
    main()