import os

import dataiku
from backend.database.store import SQLiteStore
from dataiku.customrecipe import get_output_names_for_role, get_recipe_config


def main():
    config = get_recipe_config()

    project_key = config.get("project_key")
    webapp_id = config.get("webapp_id")
    table_exports = config.get("table_exports", [])
    decompress_blobs = config.get("decompress_blobs", True)

    selected_tables = [export["table_name"] for export in table_exports]

    if not project_key:
        raise ValueError("Project key is required")
    if not webapp_id:
        raise ValueError("Webapp ID is required")
    if not table_exports or len(selected_tables) == 0:
        raise ValueError("At least one table must be added to 'Tables to Export'")

    duplicates = [table for table in selected_tables if selected_tables.count(table) > 1]
    if duplicates:
        unique_duplicates = list(set(duplicates))
        raise ValueError(
            f"Duplicate table(s) found: {', '.join(unique_duplicates)}. Each table can only be exported once."
        )

    dip_home = os.environ.get("DIP_HOME")
    if not dip_home:
        raise ValueError("DIP_HOME environment variable not found")

    workload_folder = os.path.join(dip_home, "workload-folders", "webapps", project_key, webapp_id)

    if not os.path.exists(workload_folder):
        raise ValueError(f"Webapp workload folder not found: {workload_folder}")

    db_path = os.path.join(workload_folder, "data_store.db")
    if not os.path.exists(db_path):
        raise ValueError(f"Database file not found: {db_path}")

    store = SQLiteStore(db_path, max_artifacts_size_mb=None)

    output_names = get_output_names_for_role("output")

    if len(output_names) != len(selected_tables):
        table_list = "\n".join(f"  Table {i + 1}: {table}" for i, table in enumerate(selected_tables))
        output_list = "\n".join(f"  {i + 1}. {name}" for i, name in enumerate(output_names))

        error_msg = f"""
║  OUTPUT DATASET COUNT MISMATCH  ║

You configured {len(selected_tables)} table export(s):
{table_list}

But have {len(output_names)} output dataset(s) in the Flow:
{output_list}

HOW TO FIX:
1. Edit this recipe
2. In the output section, add/remove datasets to match the number of table exports
3. Order matters: Table 1 → Output 1, Table 2 → Output 2, etc.

Required: {len(selected_tables)} output dataset(s)
Current:  {len(output_names)} output dataset(s)
"""
        raise ValueError(error_msg)

    print("=" * 50)
    print("EXPORT PLAN")
    print("=" * 50)
    print(f"Source: {project_key}/{webapp_id}")
    print(f"Database: {db_path}")
    print(f"\nTable → Dataset Mapping:")
    for i, (table, dataset) in enumerate(zip(selected_tables, output_names), 1):
        print(f"  {i}. {table:25s} → {dataset}")
    print("=" * 50)
    print()

    total_rows_exported = 0

    for table_name, output_name in zip(selected_tables, output_names):
        print(f"Exporting '{table_name}' to '{output_name}'...")

        schema_columns = store.get_table_schema(table_name)
        output_dataset = dataiku.Dataset(output_name)
        output_dataset.write_schema(schema_columns)

        row_count = 0
        with output_dataset.get_writer() as writer:
            for row in store.export_table_rows(table_name, decompress=decompress_blobs):
                row_array = [row.get(col["name"]) for col in schema_columns]
                writer.write_row_array(row_array)
                row_count += 1

        print(f"Exported {row_count:,} rows")
        total_rows_exported += row_count

    print(f"\n✓ Successfully exported {len(selected_tables)} tables ({total_rows_exported:,} total rows)")


if __name__ == "__main__":
    main()
