"""
Generate a template script for creating MLflow evaluation datasets.

This script creates a customized Python script for dataset creation,
handling both OSS MLflow and Databricks Unity Catalog scenarios.

Usage:
    # Minimum required
    python create_dataset_template.py --test-cases-file test_cases.txt

    # With custom dataset name
    python create_dataset_template.py --test-cases-file test_cases.txt --dataset-name my-eval

    # For Databricks Unity Catalog
    python create_dataset_template.py --test-cases-file test_cases.txt \\
        --catalog main --schema ml --table eval_v1
"""

import argparse
import datetime
import os
import subprocess
import sys

from utils import check_databricks_config, validate_env_vars


def parse_arguments():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(description="Generate dataset creation script")
    parser.add_argument(
        "--test-cases-file",
        required=True,
        help="File with test cases (one per line, minimum 10)",
    )
    parser.add_argument(
        "--dataset-name",
        help="Dataset name (auto-generated if not provided)",
    )
    parser.add_argument("--catalog", help="UC catalog name (for Databricks)")
    parser.add_argument("--schema", help="UC schema name (for Databricks)")
    parser.add_argument("--table", help="UC table name (for Databricks)")
    parser.add_argument("--output", default="create_evaluation_dataset.py", help="Output file name")
    return parser.parse_args()


def load_test_cases_from_file(file_path: str) -> list[str]:
    """Load test cases from file (one per line)."""
    try:
        with open(file_path) as f:
            test_cases = [line.strip() for line in f if line.strip()]

        if len(test_cases) < 10:
            print(f"✗ File has only {len(test_cases)} test cases (minimum: 10)")
            print("  Please add more test cases to the file")
            sys.exit(1)

        print(f"✓ Loaded {len(test_cases)} test cases from {file_path}")
        return test_cases

    except FileNotFoundError:
        print(f"✗ File not found: {file_path}")
        sys.exit(1)
    except Exception as e:
        print(f"✗ Error reading file: {e}")
        sys.exit(1)


def get_uc_catalogs():
    """Get available Unity Catalog catalogs."""
    try:
        code = """
from databricks import sdk
w = sdk.WorkspaceClient()
catalogs = w.catalogs.list()
for catalog in catalogs:
    print(catalog.name)
"""
        result = subprocess.run(["python", "-c", code], capture_output=True, text=True, check=True)
        return [line.strip() for line in result.stdout.strip().split("\n") if line.strip()]
    except Exception:
        return []


def get_uc_schemas(catalog: str):
    """Get available schemas in a catalog."""
    try:
        code = f"""
from databricks import sdk
w = sdk.WorkspaceClient()
schemas = w.schemas.list(catalog_name='{catalog}')
for schema in schemas:
    print(schema.name)
"""
        result = subprocess.run(["python", "-c", code], capture_output=True, text=True, check=True)
        return [line.strip() for line in result.stdout.strip().split("\n") if line.strip()]
    except Exception:
        return []


def generate_dataset_creation_code(
    tracking_uri: str,
    experiment_id: str,
    dataset_name: str,
    test_cases: list[str],
    catalog: str = None,
    schema: str = None,
    table: str = None,
) -> str:
    """Generate Python code for creating the dataset."""

    # Escape test cases for Python code
    test_cases_repr = repr(test_cases)

    if catalog and schema and table:
        # Databricks Unity Catalog version
        uc_name = f"{catalog}.{schema}.{table}"
        return f'''#!/usr/bin/env python3
"""
Create MLflow evaluation dataset.

Generated by create_dataset_template.py
"""

import mlflow
from mlflow.genai.datasets import create_dataset

# Set tracking URI
mlflow.set_tracking_uri("{tracking_uri}")

# Configuration
DATASET_NAME = "{uc_name}"
EXPERIMENT_ID = "{experiment_id}"

# Test cases
TEST_CASES = {test_cases_repr}

print("=" * 60)
print("Creating MLflow Evaluation Dataset")
print("=" * 60)
print(f"Dataset: {{DATASET_NAME}}")
print(f"Test cases: {{len(TEST_CASES)}}")
print()

# Create dataset
print("Creating dataset...")
try:
    dataset = create_dataset(
        name=DATASET_NAME,
        source={{
            "inputs": [
                {{"query": query}}
                for query in TEST_CASES
            ]
        }},
        targets=[
            {{"expected_output": "TODO: Add expected output"}}
            for _ in TEST_CASES
        ],
        experiment_id=[EXPERIMENT_ID]
    )

    print(f"✓ Dataset created: {{DATASET_NAME}}")
    print(f"  Location: Unity Catalog table")
    print(f"  Queries: {{len(TEST_CASES)}}")
    print()
    print("=" * 60)
    print("Next Steps")
    print("=" * 60)
    print()
    print("1. Verify dataset in Databricks Unity Catalog")
    print(f"2. Use in evaluation: python scripts/run_evaluation_template.py --dataset-name {uc_name}")
    print()

except Exception as e:
    print(f"✗ Error creating dataset: {{e}}")
    import traceback
    traceback.print_exc()
    exit(1)
'''
    else:
        # OSS MLflow version (non-UC)
        return f'''#!/usr/bin/env python3
"""
Create MLflow evaluation dataset.

Generated by create_dataset_template.py
"""

import mlflow
from mlflow.genai.datasets import create_dataset

# Set tracking URI
mlflow.set_tracking_uri("{tracking_uri}")

# Configuration
DATASET_NAME = "{dataset_name}"
EXPERIMENT_ID = "{experiment_id}"

# Test cases
TEST_CASES = {test_cases_repr}

print("=" * 60)
print("Creating MLflow Evaluation Dataset")
print("=" * 60)
print(f"Dataset: {{DATASET_NAME}}")
print(f"Test cases: {{len(TEST_CASES)}}")
print()

# Create dataset
print("Creating dataset...")
try:
    dataset = create_dataset(
        name=DATASET_NAME,
        source={{
            "inputs": [
                {{"query": query}}
                for query in TEST_CASES
            ]
        }},
        targets=[
            {{"expected_output": "TODO: Add expected output"}}
            for _ in TEST_CASES
        ],
        experiment_id=[EXPERIMENT_ID]
    )

    print(f"✓ Dataset created: {{DATASET_NAME}}")
    print(f"  Queries: {{len(TEST_CASES)}}")
    print()
    print("=" * 60)
    print("Next Steps")
    print("=" * 60)
    print()
    print("1. Verify dataset: python scripts/list_datasets.py")
    print(f"2. Use in evaluation: python scripts/run_evaluation_template.py --dataset-name {{DATASET_NAME}}")
    print()

except Exception as e:
    print(f"✗ Error creating dataset: {{e}}")
    import traceback
    traceback.print_exc()
    exit(1)
'''


def main():
    """Main workflow."""
    args = parse_arguments()

    print("=" * 60)
    print("MLflow Dataset Creation Template Generator")
    print("=" * 60)
    print()

    # Check environment
    errors = validate_env_vars()
    if errors:
        print("✗ Environment validation failed:")
        for error in errors:
            print(f"  - {error}")
        print("\nRun scripts/setup_mlflow.py first")
        sys.exit(1)

    tracking_uri = os.getenv("MLFLOW_TRACKING_URI")
    experiment_id = os.getenv("MLFLOW_EXPERIMENT_ID")

    print(f"Tracking URI: {tracking_uri}")
    print(f"Experiment ID: {experiment_id}")
    print()

    # Load test cases
    print("Loading test cases...")
    test_cases = load_test_cases_from_file(args.test_cases_file)
    print()

    # Check if Databricks
    is_databricks, profile = check_databricks_config()

    # Handle dataset configuration
    if is_databricks:
        print("✓ Detected Databricks configuration")
        print()

        # Check if UC args provided
        if args.catalog and args.schema and args.table:
            catalog = args.catalog
            schema = args.schema
            table = args.table
            print(f"Using Unity Catalog: {catalog}.{schema}.{table}")
        elif args.catalog or args.schema or args.table:
            print("✗ For Databricks UC, all three args are required:")
            print("  --catalog, --schema, --table")
            sys.exit(1)
        else:
            # Try to get catalogs
            print("Fetching available catalogs...")
            catalogs = get_uc_catalogs()

            if not catalogs:
                print("✗ Could not fetch catalogs")
                print("  Please specify manually:")
                print("    --catalog <catalog> --schema <schema> --table <table>")
                sys.exit(1)

            # Auto-select first catalog
            catalog = catalogs[0]
            print(f"✓ Found {len(catalogs)} catalog(s), using: {catalog}")

            # Get schemas
            print(f"Fetching schemas in {catalog}...")
            schemas = get_uc_schemas(catalog)

            if not schemas:
                print("✗ Could not fetch schemas")
                print("  Please specify manually:")
                print("    --catalog <catalog> --schema <schema> --table <table>")
                sys.exit(1)

            # Auto-select first schema
            schema = schemas[0]
            print(f"✓ Found {len(schemas)} schema(s), using: {schema}")

            # Auto-generate table name
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            table = f"eval_{timestamp}"
            print(f"✓ Auto-generated table name: {table}")

        print()
        dataset_name = None  # UC uses catalog.schema.table
    else:
        # Non-Databricks: use simple dataset name
        catalog = schema = table = None

        if args.dataset_name:
            dataset_name = args.dataset_name
            print(f"Using dataset name: {dataset_name}")
        else:
            # Auto-generate dataset name
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            dataset_name = f"agent-eval-{timestamp}"
            print(f"✓ Auto-generated dataset name: {dataset_name}")
        print()

    # Generate code
    print("=" * 60)
    print("Generating Dataset Creation Script")
    print("=" * 60)

    code = generate_dataset_creation_code(
        tracking_uri, experiment_id, dataset_name, test_cases, catalog, schema, table
    )

    # Write to file
    output_file = args.output
    with open(output_file, "w") as f:
        f.write(code)

    print(f"\n✓ Script generated: {output_file}")
    print()

    # Make executable
    try:
        os.chmod(output_file, 0o755)
        print(f"✓ Made executable: chmod +x {output_file}")
    except Exception:
        pass

    print()
    print("=" * 60)
    print("Next Steps")
    print("=" * 60)
    print()
    print(f"1. Review the generated script: {output_file}")
    print(f"2. Execute it: python {output_file}")
    print("3. Verify dataset was created: python scripts/list_datasets.py")
    print()
    print("=" * 60)


if __name__ == "__main__":
    main()
