"""Generate and save query_context for API-created Superset charts."""
import json

app = __import__("superset.app", fromlist=["create_app"]).create_app()

with app.app_context():
    from flask import g
    from superset.extensions import db
    from superset.models.slice import Slice
    from superset.models.core import Database
    from superset.charts.schemas import ChartDataQueryContextSchema
    from superset import security_manager

    admin = security_manager.find_user(username="admin")
    g.user = admin

    charts = db.session.query(Slice).order_by(Slice.id).all()
    for sl in charts:
        try:
            fd = sl.form_data
            metric = fd.get("metric")
            metrics = fd.get("metrics") or ([metric] if metric else [])
            if not metrics:
                metrics = [
                    {
                        "expressionType": "SQL",
                        "sqlExpression": "COUNT(*)",
                        "label": "COUNT(*)",
                    }
                ]

            groupby = fd.get("groupby") or []
            payload = {
                "datasource": {
                    "id": sl.datasource_id,
                    "type": sl.datasource_type,
                },
                "force": False,
                "queries": [
                    {
                        "filters": [],
                        "extras": {"having": "", "where": ""},
                        "applied_time_extras": {},
                        "columns": groupby if isinstance(groupby, list) else [],
                        "metrics": metrics,
                        "orderby": [],
                        "annotation_layers": [],
                        "row_limit": int(fd.get("row_limit") or 1000),
                        "series_limit": 0,
                        "order_desc": True,
                        "url_params": {},
                        "custom_params": {},
                        "custom_form_data": {},
                    }
                ],
                "form_data": fd,
                "result_format": "json",
                "result_type": "full",
            }
            qc = ChartDataQueryContextSchema().load(payload)
            ctx = qc.cache_values if hasattr(qc, "cache_values") else None
            if ctx is None:
                # fallback: store factory input dict
                from superset.common.query_context_factory import QueryContextFactory

                factory = QueryContextFactory()
                ctx = {
                    "datasource": {
                        "id": sl.datasource_id,
                        "type": sl.datasource_type,
                    },
                    "force": False,
                    "queries": payload["queries"],
                    "form_data": fd,
                    "result_format": "json",
                    "result_type": "full",
                }
            sl.query_context = json.dumps(ctx) if isinstance(ctx, dict) else json.dumps(payload)
            sl.query_context_generation = True
            db.session.add(sl)
            print("OK", sl.id, sl.slice_name[:50])
        except Exception as e:
            print("ERR", sl.id, sl.slice_name[:40], e)
    db.session.commit()
    print("committed")