#!/usr/bin/env python3
"""Generate SQL migration to change the pgvector embedding dimension.

Usage:
    python sql/generate_dimension_migration.py 1024                       # postgres (default)
    python sql/generate_dimension_migration.py 1024 --schema supabase     # Supabase Cloud
    python sql/generate_dimension_migration.py 1024 --schema selfhost     # Supabase Docker+RLS
    python sql/generate_dimension_migration.py 512 > sql/migrations/012_resize_512.sql

Schemas:
    postgres  — Neon / vanilla PostgreSQL (no RLS, default)
    supabase  — Supabase Cloud (extensions schema, RLS)
    selfhost  — Self-hosted Supabase Docker (RLS)

Reads the appropriate schema file, replaces the vector dimension, and outputs
a complete migration script you can run via psql or paste into the SQL editor.
"""

import argparse
import re
import sys
from pathlib import Path

SQL_DIR = Path(__file__).parent
SCHEMA_FILES = {
    "selfhost": SQL_DIR / "schema_selfhost_supabase.sql",
    "postgres": SQL_DIR / "schema_postgres.sql",
    "supabase": SQL_DIR / "schema.sql",
}


def extract_functions(schema_text: str) -> list[str]:
    """Extract all CREATE OR REPLACE FUNCTION blocks from schema SQL."""
    # Match CREATE OR REPLACE FUNCTION ... $$; blocks
    pattern = re.compile(
        r"(CREATE OR REPLACE FUNCTION\s+\w+\s*\(.*?\)\s*"
        r"RETURNS\s+.*?"
        r"\$\$;)",
        re.DOTALL | re.IGNORECASE,
    )
    return pattern.findall(schema_text)


def has_vector_param(function_sql: str) -> bool:
    """Check if a function has a vector type in its parameters."""
    return bool(re.search(r"vector\(\d+\)", function_sql, re.IGNORECASE))


def generate_migration(target_dim: int, schema_type: str) -> str:
    schema_file = SCHEMA_FILES[schema_type]
    if not schema_file.exists():
        print(f"Schema file not found: {schema_file}", file=sys.stderr)
        sys.exit(1)

    schema_text = schema_file.read_text()

    # Detect current dimension from the table definition
    match = re.search(r"vector\((\d+)\)", schema_text)
    if not match:
        print("Could not detect current vector dimension in schema", file=sys.stderr)
        sys.exit(1)
    current_dim = int(match.group(1))

    if current_dim == target_dim:
        print(f"Schema already uses vector({target_dim})", file=sys.stderr)
        sys.exit(1)

    # Extract functions that reference vector types
    functions = extract_functions(schema_text)
    vector_functions = [f for f in functions if has_vector_param(f)]

    # Build migration SQL
    lines = [
        f"-- Migration: vector({current_dim}) -> vector({target_dim})",
        f"-- Generated by: python sql/generate_dimension_migration.py {target_dim}",
        f"-- Schema: {schema_type}",
        "",
        "-- Step 1: Null out existing vectors (pgvector can't cast between dimensions)",
        "UPDATE memories SET embedding = NULL;",
        "",
        "-- Step 2: Alter the column type",
        f"ALTER TABLE memories ALTER COLUMN embedding TYPE "
        f"{'extensions.' if schema_type == 'supabase' else ''}vector({target_dim});",
        "",
        "-- Step 3: Update function signatures",
    ]

    for func in vector_functions:
        # Get the function name for a comment
        name_match = re.search(r"FUNCTION\s+(\w+)", func, re.IGNORECASE)
        func_name = name_match.group(1) if name_match else "unknown"

        # Replace dimension
        updated = re.sub(r"vector\(\d+\)", f"vector({target_dim})", func, flags=re.IGNORECASE)
        # For supabase schema, also handle extensions.vector
        if schema_type == "supabase":
            updated = re.sub(
                r"extensions\.vector\(\d+\)",
                f"extensions.vector({target_dim})",
                updated,
                flags=re.IGNORECASE,
            )

        lines.append("")
        lines.append(f"-- {func_name}")
        lines.append(updated)

    lines.extend(
        [
            "",
            "-- Step 4: Restart the MCP server and run re_embed_all",
            f"-- All embeddings are NULL after this migration."
            f" Run re_embed_all to regenerate at {target_dim} dimensions.",
        ]
    )

    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(
        description="Generate SQL migration to change pgvector embedding dimension"
    )
    parser.add_argument("dimension", type=int, help="Target vector dimension (e.g. 512, 768, 1024)")
    parser.add_argument(
        "--schema",
        choices=["supabase", "selfhost", "postgres"],
        default="postgres",
        help="Schema: supabase (cloud), selfhost (Docker+RLS), postgres (Neon). Default: postgres",
    )
    args = parser.parse_args()

    if args.dimension < 1 or args.dimension > 4096:
        print("Dimension must be between 1 and 4096", file=sys.stderr)
        sys.exit(1)

    print(generate_migration(args.dimension, args.schema))


if __name__ == "__main__":
    main()
