from __future__ import annotations

import os
from contextlib import contextmanager

import dataiku
from dataiku.customrecipe import get_output_names_for_role, get_recipe_config


def load_recipe_inputs():
    config = get_recipe_config()
    table_exports = config.get("table_exports", [])

    return {
        "config": config,
        "project_key": config.get("project_key"),
        "webapp_id": config.get("webapp_id"),
        "table_exports": table_exports,
        "selected_tables": [export["table_name"] for export in table_exports],
        "decompress_blobs": config.get("decompress_blobs", True),
    }


def validate_recipe_selection(project_key, webapp_id, selected_tables):
    if not project_key:
        raise ValueError("Project key is required")
    if not webapp_id:
        raise ValueError("Webapp ID is required")
    if not selected_tables:
        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."
        )


def load_output_names(selected_tables):
    output_names = get_output_names_for_role("output")
    if len(output_names) == len(selected_tables):
        return output_names

    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))

    raise ValueError(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)
""")


@contextmanager
def temporary_env(**updates):
    previous_values = {key: os.environ.get(key) for key in updates}
    try:
        for key, value in updates.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value
        yield
    finally:
        for key, previous_value in previous_values.items():
            if previous_value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = previous_value


def print_export_plan(project_key, webapp_id, selected_tables, output_names, webapp_name=None, webapp_label=None):
    print("=" * 50)
    print("EXPORT PLAN")
    print("=" * 50)
    print(f"Source: {project_key}/{webapp_id}")
    if webapp_name and webapp_label:
        print(f"Webapp: {webapp_name} ({webapp_label})")
    print("\nTable → Dataset Mapping:")
    for i, (table, dataset) in enumerate(zip(selected_tables, output_names), 1):
        print(f"  {i}. {table:25s} → {dataset}")
    print("=" * 50)
    print()


def export_selected_tables(
    selected_tables,
    output_names,
    db_url,
    workload_folder,
    decompress_blobs,
    selected_table_display_names=None,
):
    from backend.services.export_service import ExportService

    export_service = ExportService(db_url, workload_folder)
    total_rows_exported = 0

    if selected_table_display_names and len(selected_table_display_names) != len(selected_tables):
        raise ValueError("selected_table_display_names must match selected_tables length")

    for i, (table_name, output_name) in enumerate(zip(selected_tables, output_names)):
        display_name = selected_table_display_names[i] if selected_table_display_names else table_name

        print(f"Exporting '{display_name}' to '{output_name}'...")

        schema_columns = export_service.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 export_service.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)")
