"""
Airflow 3 DAG: Migration and Sync Pipeline for AI Law Bot
Extracts law text from local MongoDB and upserts into the new PostgreSQL database.
Embeddings will be automatically handled by `pgai` vectorizer within Postgres.
"""
from datetime import datetime, timedelta
from airflow.decorators import dag, task
from airflow.providers.mongo.hooks.mongo import MongoHook
from airflow.providers.postgres.hooks.postgres import PostgresHook
import pandas as pd

# Define Default Arguments
default_args = {
    'owner': 'ai_team',
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
}
@dag(
    dag_id='ai_law_sync_to_postgres',
    default_args=default_args,
    description='Syncs law documents from MongoDB to PostgreSQL for pgai vectorization',
    schedule_interval='@daily',
    start_date=datetime(2026, 4, 1),
    catchup=False,
    tags=['law_bot', 'ai', 'migration'],
)
def ai_law_postgres_sync():
    @task()
    def initialize_postgres_schema():
        """
        Ensures the target PostgreSQL table exists.
        Note: The pgai vectorizer should be configured independently on the database level.
        """
        pg_hook = PostgresHook(postgres_conn_id='canifa_postgres_138')
        schema_query = """
        CREATE TABLE IF NOT EXISTS law_documents (
            id VARCHAR(255) PRIMARY KEY,
            title TEXT,
            content TEXT NOT NULL,
            category VARCHAR(100),
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
        """
        pg_hook.run(schema_query)
        print("Schema initialized.")
    @task(multiple_outputs=True)
    def extract_from_mongodb() -> dict:
        """
        Extracts law documents from MongoDB (`ai_law` database).
        """
        mongo_hook = MongoHook(mongo_conn_id='local_mongo_law')
        client = mongo_hook.get_conn()
        db = client.ai_law
        collection = db.law_records

        # In a real scenario, you'd fetch based on a last_sync_date (incremental loading).
        # For this template, we fetch a batch of records.
        records = list(collection.find({}, {'_id': 1, 'title': 1, 'content': 1, 'category': 1}).limit(10000))

        # Convert _id ObjectId to string format
        for record in records:
            record['id'] = str(record.pop('_id'))

        print(f"Extracted {len(records)} documents from MongoDB.")
        return {'records': records}
    @task()
    def load_to_postgres(extracted_data: dict):
        """
        Upserts the extracted data into PostgreSQL.
        """
        records = extracted_data.get('records', [])
        if not records:
            print("No records to insert. Skipping.")
            return
        # Convert to Pandas DataFrame for efficient bulk insertion
        df = pd.DataFrame(records)

        pg_hook = PostgresHook(postgres_conn_id='canifa_postgres_138')
        engine = pg_hook.get_sqlalchemy_engine()

        # Temporary table logic or direct 'to_sql' based on upsert preference.
        # This uses direct append for simplicity; for real upsert, use SQLAlchemy raw SQL
        df.to_sql('law_documents', con=engine, if_exists='append', index=False, method='multi', chunksize=1000)

        print(f"Successfully loaded {len(records)} records into PostgreSQL.")
    # DAG Execution Flow
    schema_setup = initialize_postgres_schema()
    extracted_data = extract_from_mongodb()

    schema_setup >> extracted_data
    load_to_postgres(extracted_data)
# Instantiate the DAG
dag = ai_law_postgres_sync()
