#!/usr/bin/env python3
"""Migration: Add 'strang' column to ketten and populate based on Ursprungs-Vorlage typ."""

import sqlite3
import sys
from pathlib import Path

DB_PATH = Path(__file__).parent.parent / "data" / "tracker.db"


VORLAGE_TYP_TO_STRANG = {
    "antrag": "antrag",
    "anfrage": "anfrage",
    "beschlussvorlage": "beschlussvorlage",
    "mitteilungsvorlage": "mitteilung",
}


def migrate(db_path: Path = DB_PATH):
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row

    # Check if column already exists
    cols = [r[1] for r in conn.execute("PRAGMA table_info(ketten)").fetchall()]
    if "strang" not in cols:
        print("Adding 'strang' column to ketten...")
        conn.execute("ALTER TABLE ketten ADD COLUMN strang TEXT")
        conn.commit()
        print("Column added.")
    else:
        print("Column 'strang' already exists.")

    # Populate strang based on Ursprungs-Vorlage typ
    rows = conn.execute("""
        SELECT k.id, v.typ as vorlage_typ
        FROM ketten k
        LEFT JOIN vorlagen v ON k.ursprung_id = v.id
        WHERE k.strang IS NULL
    """).fetchall()

    if not rows:
        print("All ketten already have strang set.")
        conn.close()
        return

    print(f"Updating {len(rows)} ketten...")
    counts = {}
    for r in rows:
        vorlage_typ = r["vorlage_typ"] or ""
        strang = VORLAGE_TYP_TO_STRANG.get(vorlage_typ, "sonstig")
        conn.execute("UPDATE ketten SET strang = ? WHERE id = ?", (strang, r["id"]))
        counts[strang] = counts.get(strang, 0) + 1

    conn.commit()
    conn.close()

    print("Done. Counts per strang:")
    for strang, count in sorted(counts.items()):
        print(f"  {strang}: {count}")


if __name__ == "__main__":
    db = Path(sys.argv[1]) if len(sys.argv) > 1 else DB_PATH
    migrate(db)