#!/usr/bin/env python3
"""Split models.py into domain-specific modules.

Target modules:
  - migrations.py: all migrate_* functions
  - helpers.py: utility functions (geo, UA, encryption helpers)
  - analytics.py: analytics, stats, trends, dropoff functions
  - documents.py: document templates, invoices, schedules, QR, email HTML
  - email_campaigns.py: campaigns, recipients, tracking, A/B tests, send quota, reminders

Approach:
  1. Parse models.py to find function boundaries
  2. Categorize each function
  3. Extract code blocks to new files
  4. Update models.py to import + re-export everything
"""

import re
import sys

MODEL_PATH = "app/models.py"

# Header to copy to each new module (shared imports and constants)
SHARED_HEADER = '''\
"""Auto-extracted from models.py — do not edit manually."""
import sqlite3
import os
import secrets
import uuid
from datetime import datetime, timezone
import bcrypt
import json

from app.crypto import (
    encrypt_value, try_decrypt, key_is_configured,
    encrypt_user_value, decrypt_user_value, try_decrypt_user_value,
    try_decrypt_user_submission, encrypt_user_submission,
    encrypt_entity, decrypt_entity, try_decrypt_entity,
    hash_value, is_encrypted,
)

DB_PATH = os.environ.get("RELAY_DB_PATH", "/app/data/relay.db")

'''

# Functions that belong in each module
MIGRATION_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',
}

HELPER_FUNCS = {
    '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',
}

ANALYTICS_FUNCS = {
    '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',
}

DOCUMENT_FUNCS = {
    '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',
}

EMAIL_FUNCS = {
    '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',
}


def parse_functions(filepath):
    """Parse function definitions and their line ranges."""
    with open(filepath, 'r') as f:
        lines = f.readlines()

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

    functions = []
    for i, line in enumerate(lines):
        m = func_pattern.match(line)
        if m:
            functions.append({
                'line': i,
                'async': m.group(1) is not None,
                'name': m.group(2),
            })

    # Set end lines (one before next function)
    for i, func in enumerate(functions):
        if i + 1 < len(functions):
            func['end'] = functions[i + 1]['line']
        else:
            func['end'] = len(lines)

    return lines, functions


def extract_blocks(lines, func_names):
    """Extract code blocks for a set of function names, preserving order."""
    results = []
    func_pattern = re.compile(r'^(async\s+)?def\s+(' + '|'.join(func_names) + r')\s*\(')

    i = 0
    while i < len(lines):
        m = func_pattern.match(lines[i])
        if m:
            func_name = m.group(2)
            start = i
            # Find end of this function
            j = i + 1
            while j < len(lines):
                if lines[j].startswith('def ') or lines[j].startswith('class '):
                    break
                j += 1
            results.append((func_name, lines[start:j]))
            i = j
        else:
            i += 1

    return results


def main():
    lines, functions = parse_functions(MODEL_PATH)

    # Verify all target functions exist
    all_targets = MIGRATION_FUNCS | HELPER_FUNCS | ANALYTICS_FUNCS | DOCUMENT_FUNCS | EMAIL_FUNCS
    found_names = {f['name'] for f in functions}
    missing = all_targets - found_names
    if missing:
        print(f"WARNING: Functions not found in models.py: {missing}")

    # Create migrations.py
    print("Creating app/migrations.py ...")
    block = extract_blocks(lines, MIGRATION_FUNCS)
    write_module("app/migrations.py", block, lines, "migrations")

    # Create helpers.py
    print("Creating app/helpers.py ...")
    block = extract_blocks(lines, HELPER_FUNCS)
    write_module("app/helpers.py", block, lines, "helpers")

    # Create analytics.py
    print("Creating app/analytics.py ...")
    block = extract_blocks(lines, ANALYTICS_FUNCS)
    write_module("app/analytics.py", block, lines, "analytics")

    # Create documents.py (new file, rename existing services/documents.py later)
    print("Creating app/model_documents.py ...")
    block = extract_blocks(lines, DOCUMENT_FUNCS)
    write_module("app/model_documents.py", block, lines, "model_documents")

    # Create email_campaigns.py
    print("Creating app/model_email_campaigns.py ...")
    block = extract_blocks(lines, EMAIL_FUNCS)
    write_module("app/model_email_campaigns.py", block, lines, "model_email_campaigns")

    # Now generate the re-export stubs for models.py
    all_migrated = MIGRATION_FUNCS | HELPER_FUNCS | ANALYTICS_FUNCS | DOCUMENT_FUNCS | EMAIL_FUNCS

    # Build import statements
    imports = []
    imports.append("from app.migrations import (")
    for name in sorted(MIGRATION_FUNCS):
        imports.append(f"    {name},")
    imports.append(")")

    imports.append("from app.helpers import (")
    for name in sorted(HELPER_FUNCS):
        imports.append(f"    {name},")
    imports.append(")")

    imports.append("from app.analytics import (")
    for name in sorted(ANALYTICS_FUNCS):
        imports.append(f"    {name},")
    imports.append(")")

    imports.append("from app.model_documents import (")
    for name in sorted(DOCUMENT_FUNCS):
        imports.append(f"    {name},")
    imports.append(")")

    imports.append("from app.model_email_campaigns import (")
    for name in sorted(EMAIL_FUNCS):
        imports.append(f"    {name},")
    imports.append(")")

    print(f"\nFunctions to migrate: {len(all_migrated)}")
    print(f"Functions remaining in models.py: {len(found_names - all_migrated)}")

    # Write the import block to a temp file for manual insertion
    with open("/tmp/models_imports.txt", 'w') as f:
        f.write('\n'.join(imports) + '\n')

    print("\nImport block written to /tmp/models_imports.txt")
    print("New modules created. Review before replacing models.py.")


def write_module(filepath, blocks, all_lines, module_name):
    """Write a module file with header + extracted function blocks."""
    with open(filepath, 'w') as f:
        f.write(SHARED_HEADER)
        f.write('\n\n')

        for func_name, func_lines in blocks:
            for line in func_lines:
                f.write(line)
            f.write('\n')

    print(f"  → {filepath} ({len(blocks)} functions)")


if __name__ == '__main__':
    main()