"""
Sync dbt marts → Typesense search index.

Reads dim_school + dim_location + latest fact data from PostgreSQL marts,
then upserts into a Typesense collection with zero-downtime alias swapping.

Usage:
    python sync_typesense.py [--typesense-url http://localhost:8108] [--api-key xyz]
"""

from __future__ import annotations

import argparse
import os
import sys
import time

import psycopg2
import psycopg2.extras
import typesense

from gias_codes import PHASE_OF_EDUCATION, RELIGIOUS_CHARACTER, SCHOOL_TYPE, translate

COLLECTION_SCHEMA = {
    "fields": [
        {"name": "urn", "type": "int32"},
        {"name": "school_name", "type": "string"},
        {"name": "phase", "type": "string", "facet": True},
        {"name": "school_type", "type": "string", "facet": True},
        {"name": "local_authority", "type": "string", "facet": True},
        {"name": "religious_character", "type": "string", "facet": True, "optional": True},
        {"name": "ofsted_rating", "type": "string", "facet": True, "optional": True},
        {"name": "postcode", "type": "string"},
        {"name": "location", "type": "geopoint", "optional": True},
        {"name": "headteacher_name", "type": "string", "optional": True},
        {"name": "rwm_expected_pct", "type": "float", "optional": True},
        {"name": "progress_8_score", "type": "float", "optional": True},
        {"name": "total_pupils", "type": "int32"},
    ],
    "default_sorting_field": "total_pupils",
}

OFSTED_LABELS = {1: "Outstanding", 2: "Good", 3: "Requires Improvement", 4: "Inadequate"}

QUERY_BASE = """
    SELECT
        s.urn,
        s.school_name,
        s.phase_code,
        s.school_type_code,
        l.local_authority_name as local_authority,
        s.religious_character_code,
        s.ofsted_grade,
        l.postcode,
        s.headteacher_name,
        s.total_pupils,
        l.latitude as lat,
        l.longitude as lng
    FROM marts.dim_school s
    LEFT JOIN marts.dim_location l ON s.urn = l.urn
"""

QUERY_PERFORMANCE_JOIN = """
    LEFT JOIN LATERAL (
        SELECT rwm_expected_pct, progress_8_score
        FROM marts.fact_performance
        WHERE urn = s.urn
        ORDER BY year DESC
        LIMIT 1
    ) p ON true
"""


def get_db_connection():
    return psycopg2.connect(
        host=os.environ.get("PG_HOST", "localhost"),
        port=os.environ.get("PG_PORT", "5432"),
        user=os.environ.get("PG_USER", "postgres"),
        password=os.environ.get("PG_PASSWORD", "postgres"),
        dbname=os.environ.get("PG_DATABASE", "school_compare"),
    )


def build_document(row: dict) -> dict:
    """Convert a DB row to a Typesense document."""
    doc = {
        "id": str(row["urn"]),
        "urn": row["urn"],
        "school_name": row["school_name"] or "",
        "phase": translate(row["phase_code"], PHASE_OF_EDUCATION) or "",
        "school_type": translate(row["school_type_code"], SCHOOL_TYPE) or "",
        "local_authority": row["local_authority"] or "",
        "postcode": row["postcode"] or "",
    }

    religious_character = translate(row.get("religious_character_code"), RELIGIOUS_CHARACTER)
    if religious_character:
        doc["religious_character"] = religious_character
    if row.get("ofsted_grade"):
        doc["ofsted_rating"] = OFSTED_LABELS.get(row["ofsted_grade"], "")
    if row.get("headteacher_name"):
        doc["headteacher_name"] = row["headteacher_name"]
    doc["total_pupils"] = row["total_pupils"] or 0
    if row.get("rwm_expected_pct") is not None:
        doc["rwm_expected_pct"] = float(row["rwm_expected_pct"])
    if row.get("progress_8_score") is not None:
        doc["progress_8_score"] = float(row["progress_8_score"])

    if row.get("lat") is not None and row.get("lng") is not None:
        doc["location"] = [float(row["lat"]), float(row["lng"])]

    return doc


def sync(typesense_url: str, api_key: str):
    client = typesense.Client({
        "nodes": [{"host": typesense_url.split("//")[-1].split(":")[0],
                    "port": typesense_url.split(":")[-1],
                    "protocol": "http"}],
        "api_key": api_key,
        "connection_timeout_seconds": 10,
    })

    # Create timestamped collection for zero-downtime swap
    ts = int(time.time())
    collection_name = f"schools_{ts}"

    print(f"Creating collection: {collection_name}")
    schema = {**COLLECTION_SCHEMA, "name": collection_name}
    client.collections.create(schema)

    # Fetch data from marts — join fact_performance if it exists
    conn = get_db_connection()
    with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
        # Check whether the merged fact table exists
        cur.execute("""
            SELECT table_name FROM information_schema.tables
            WHERE table_schema = 'marts' AND table_name = 'fact_performance'
        """)
        has_fact_performance = cur.fetchone() is not None

        query = QUERY_BASE
        if has_fact_performance:
            query = query.replace(
                "l.longitude as lng",
                "l.longitude as lng,\n        p.rwm_expected_pct,\n        p.progress_8_score",
            )
            query += QUERY_PERFORMANCE_JOIN

        cur.execute(query)
        rows = cur.fetchall()
    conn.close()

    print(f"Indexing {len(rows)} schools...")

    # Batch import
    batch_size = 500
    for i in range(0, len(rows), batch_size):
        batch = [build_document(r) for r in rows[i : i + batch_size]]
        client.collections[collection_name].documents.import_(batch, {"action": "upsert"})
        print(f"  Indexed {min(i + batch_size, len(rows))}/{len(rows)}")

    # Swap alias
    print("Swapping alias 'schools' → new collection")
    try:
        client.aliases.upsert("schools", {"collection_name": collection_name})
    except Exception:
        # If alias doesn't exist yet, create it
        client.aliases.upsert("schools", {"collection_name": collection_name})

    print("Done.")


def main():
    parser = argparse.ArgumentParser(description="Sync marts to Typesense")
    parser.add_argument("--typesense-url", default=os.environ.get("TYPESENSE_URL", "http://localhost:8108"))
    parser.add_argument("--api-key", default=os.environ.get("TYPESENSE_API_KEY", ""))
    args = parser.parse_args()

    if not args.api_key:
        print("Error: --api-key or TYPESENSE_API_KEY required", file=sys.stderr)
        sys.exit(1)

    sync(args.typesense_url, args.api_key)


if __name__ == "__main__":
    main()
