#!/usr/bin/env python3
"""Update models.py to remove migrated functions and add re-exports.

Reads original models.py, strips out functions that moved to other modules,
adds import/re-export statements at top, preserves everything else intact.
"""

import re
import sys

MODEL_PATH = "app/models.py"
BACKUP_PATH = "app/models.py.bak"

MIGRATED_FUNCS = {
    'migrate_add_field_config', 'migrate_add_submission_data',
    'migrate_add_user_id_to_sites', 'migrate_add_usage',
    'migrate_add_abuse_protection', 'migrate_add_rate_limit',
    'migrate_add_analytics', 'migrate_add_email_verification',
    'migrate_add_indexes', 'migrate_add_api_keys',
    'migrate_encrypt_existing_data', 'migrate_rescope_to_user_keys',
    'migrate_rotate_encryption_keys', 'migrate_normalize_field_configs',
    'migrate_add_form_templates', 'migrate_add_teams',
    'migrate_add_form_analytics', 'migrate_add_ai_builder',
    'migrate_add_webhook_destinations', 'migrate_add_form_versions',
    'migrate_add_form_sessions', 'migrate_add_form_variants',
    'migrate_add_branding', 'migrate_add_webhook_retry_locking',
    'migrate_add_document_templates', 'migrate_encrypt_user_pii',
    'migrate_add_invite_expiration', 'migrate_add_user_id_to_doc_templates',
    'migrate_add_user_id_to_invoice_schedules', 'migrate_add_email_phase3',
    'migrate_add_email_phase4', 'migrate_add_email_tracking',
    'migrate_add_email_campaigns', 'migrate_add_form_actions',
    'migrate_add_user_referral_cols',
    'resolve_geo', 'parse_user_agent', '_is_encrypted_value',
    '_table_exists', '_current_month', '_decrypt_invoice_schedule',
    '_decrypt_webhook_dest', '_decrypt_webhook_log', '_decrypt_team_member',
    '_decrypt_user_row', '_encrypt_submission_field',
    '_decrypt_submission_fields', '_get_site_owner_password_hash',
    '_check_spam_content',
    'site_session_stats', 'site_field_dropoff', 'site_hourly_trends',
    'ab_test_stats', 'get_site_analytics', 'get_site_analytics_summary',
    'site_geo_stats', 'site_device_stats', 'site_impression_stats',
    'get_user_analytics', 'get_team_analytics', 'get_admin_stats',
    'get_spam_submissions', 'dismiss_spam',
    'create_document_template', 'get_document_template',
    'list_document_templates', 'get_auto_generate_template',
    'update_document_template', 'delete_document_template',
    'increment_document_download', 'store_document', 'get_document',
    'list_documents', 'update_document_status', 'update_document_due_date',
    'create_invoice_schedule', 'get_invoice_schedule',
    'list_invoice_schedules', 'update_invoice_schedule',
    'delete_invoice_schedule', 'get_due_schedules',
    'advance_schedule_next_run', 'generate_qr_code',
    'generate_email_html', 'update_site_branding',
    'generate_invoice_number', 'update_next_invoice_number',
    'list_user_documents', 'count_user_documents',
    'list_documents_by_email', 'update_document_stripe_payment',
    'create_campaign', 'get_campaign', 'get_campaign_by_id',
    'list_campaigns', 'add_campaign_recipients',
    'update_campaign_status', 'delete_campaign',
    'get_campaign_recipients', 'get_pending_recipients',
    'update_recipient_status', 'get_recipient_by_message_id',
    'get_email_settings', 'get_recipient_by_token',
    'record_open', 'record_click', 'record_bounce',
    'get_campaign_recipients_by_email', 'get_campaign_analytics',
    'upsert_email_settings', 'create_reminder', 'get_reminder',
    'get_reminders', 'cancel_reminder',
    'schedule_reminders_for_campaign', 'get_pending_reminders',
    'get_reminder_eligible_recipients', 'mark_reminder_sent',
    'check_send_quota', 'record_send', 'cleanup_old_send_log',
    'create_ab_test', 'get_ab_variant', 'get_ab_variants',
    'get_ab_test_info', 'declare_ab_winner',
    'assign_recipients_to_variants', 'send_ab_winner_to_remaining',
    'check_pending_ab_tests', 'get_custom_templates',
    'get_custom_template', 'create_custom_template',
    'update_custom_template', 'delete_custom_template',
}

REEXPORT_IMPORTS = '''\
# ── Re-exports from split modules (models.py → domain modules) ──────────────
from app.migrations import (
    migrate_add_abuse_protection, migrate_add_ai_builder, migrate_add_analytics,
    migrate_add_api_keys, migrate_add_branding, migrate_add_email_campaigns,
    migrate_add_email_phase3, migrate_add_email_phase4, migrate_add_email_tracking,
    migrate_add_email_verification, migrate_add_field_config,
    migrate_add_form_actions, migrate_add_form_analytics, migrate_add_form_sessions,
    migrate_add_form_templates, migrate_add_form_versions, migrate_add_form_variants,
    migrate_add_indexes, migrate_add_invoice_schedules, migrate_add_invite_expiration,
    migrate_add_rate_limit, migrate_add_submission_data, migrate_add_teams,
    migrate_add_usage, migrate_add_user_id_to_doc_templates,
    migrate_add_user_id_to_invoice_schedules, migrate_add_user_id_to_sites,
    migrate_add_user_referral_cols, migrate_add_webhook_destinations,
    migrate_add_webhook_retry_locking, migrate_add_document_templates,
    migrate_encrypt_existing_data, migrate_encrypt_user_pii,
    migrate_normalize_field_configs, migrate_rescope_to_user_keys,
    migrate_rotate_encryption_keys,
)
from app.helpers import (
    _check_spam_content, _current_month, _decrypt_invoice_schedule,
    _decrypt_submission_fields, _decrypt_team_member, _decrypt_user_row,
    _decrypt_webhook_dest, _decrypt_webhook_log, _encrypt_submission_field,
    _get_site_owner_password_hash, _is_encrypted_value, _table_exists,
    parse_user_agent, resolve_geo,
)
from app.analytics import (
    ab_test_stats, dismiss_spam, get_admin_stats, get_site_analytics,
    get_site_analytics_summary, get_spam_submissions, get_team_analytics,
    get_user_analytics, site_device_stats, site_field_dropoff,
    site_geo_stats, site_hourly_trends, site_impression_stats,
    site_session_stats,
)
from app.model_documents import (
    advance_schedule_next_run, create_document_template,
    create_invoice_schedule, delete_document_template, delete_invoice_schedule,
    generate_email_html, generate_invoice_number, generate_qr_code,
    get_auto_generate_template, get_document, get_document_template,
    get_due_schedules, get_invoice_schedule, increment_document_download,
    list_documents, list_document_templates, list_invoice_schedules,
    list_user_documents, count_user_documents, list_documents_by_email,
    store_document, update_document_due_date, update_document_status,
    update_document_stripe_payment, update_document_template,
    update_next_invoice_number, update_site_branding,
)
from app.model_email_campaigns import (
    add_campaign_recipients, assign_recipients_to_variants,
    cancel_reminder, check_pending_ab_tests, check_send_quota,
    cleanup_old_send_log, create_ab_test, create_campaign, create_custom_template,
    create_reminder, declare_ab_winner, delete_campaign, delete_custom_template,
    get_ab_test_info, get_ab_variant, get_ab_variants, get_campaign,
    get_campaign_analytics, get_campaign_by_id, get_campaign_recipients,
    get_campaign_recipients_by_email, get_custom_template, get_email_settings,
    get_pending_recipients, get_pending_reminders, get_recipient_by_message_id,
    get_recipient_by_token, get_reminder, get_reminder_eligible_recipients,
    get_reminders, list_campaigns, mark_reminder_sent, record_bounce,
    record_click, record_open, record_send, schedule_reminders_for_campaign,
    send_ab_winner_to_remaining, update_campaign_status, update_custom_template,
    update_recipient_status, upsert_email_settings,
)
'''


def main():
    with open(MODEL_PATH, 'r') as f:
        lines = f.readlines()

    # Backup original
    with open(BACKUP_PATH, 'w') as f:
        f.writelines(lines)
    print(f"Backed up {MODEL_PATH} → {BACKUP_PATH}")

    # Find function boundaries
    func_pattern = re.compile(r'^(async\s+)?def\s+(\w+)\s*\(')

    # Determine which line ranges to remove
    keep_lines = []
    i = 0
    removed_count = 0
    while i < len(lines):
        m = func_pattern.match(lines[i])
        if m and m.group(2) in MIGRATED_FUNCS:
            # Skip this function entirely (until next function or end)
            removed_count += 1
            i += 1
            while i < len(lines):
                if lines[i].startswith('def ') or lines[i].startswith('class '):
                    break
                i += 1
            continue
        keep_lines.append(lines[i])
        i += 1

    # Find where to insert re-exports: after the existing imports, before first function
    # The first function is `def get_user_tier` at line 93 (0-indexed: 92)
    insert_point = None
    for idx, line in enumerate(keep_lines):
        if line.startswith('def ') or line.startswith('class '):
            insert_point = idx
            break

    if insert_point is None:
        print("ERROR: Could not find insertion point in models.py")
        return 1

    # Insert re-exports before first function definition
    # First, add a blank line separator before the imports
    while insert_point > 0 and keep_lines[insert_point - 1].strip() == '':
        insert_point -= 1

    keep_lines.insert(insert_point, '\n')
    keep_lines.insert(insert_point + 1, REEXPORT_IMPORTS + '\n')

    with open(MODEL_PATH, 'w') as f:
        f.writelines(keep_lines)

    original_lines = sum(1 for _ in open(BACKUP_PATH))
    new_lines = len(keep_lines)

    print(f"\nResults:")
    print(f"  Original models.py: {original_lines} lines")
    print(f"  New models.py: {new_lines} lines")
    print(f"  Functions removed: {removed_count}")
    print(f"  Lines saved: {original_lines - new_lines}")
    print(f"\nBackup: {BACKUP_PATH}")
    print("Next: run tests to verify nothing broke")

    return 0


if __name__ == '__main__':
    sys.exit(main() or 0)