"""
DLT Pipeline Runner

Simple interface to run dlt pipelines from Fabric notebooks.
Source configuration is read from YAML files or passed directly as JSON.
Secrets are resolved from Azure Key Vault using the central vault URL.
"""

import os
import sys
import subprocess
import yaml
import importlib.util
from pathlib import Path
from typing import Optional, Dict, Any

from fabric_utils import mount_lakehouse


def git_pull(repo_path: str, pat_token: str) -> bool:
    """Pull latest changes from git repo."""
    if not os.path.exists(os.path.join(repo_path, ".git")):
        print(f"  Not a git repo: {repo_path}")
        return False

    # Get repo URL and add PAT
    result = subprocess.run(
        ["git", "remote", "get-url", "origin"],
        cwd=repo_path,
        capture_output=True,
        text=True
    )

    if result.returncode == 0:
        url = result.stdout.strip()
        auth_url = url.replace("https://", f"https://{pat_token}@")

        # Update remote with auth
        subprocess.run(
            ["git", "remote", "set-url", "origin", auth_url],
            cwd=repo_path,
            capture_output=True
        )

    # Pull
    result = subprocess.run(
        ["git", "pull", "--ff-only"],
        cwd=repo_path,
        capture_output=True,
        text=True,
        timeout=60
    )

    if result.returncode == 0:
        print(f"  Pulled: {repo_path}")
        return True
    else:
        print(f"  Pull failed: {repo_path}")
        return False


def run_pipeline(
    source_name: str,
    vault_url: str,
    base_path: str = "/tmp/vd-dlt",
    destination_path: str = "/lakehouse/default/Tables",
    github_pat_secret: str = "my-github-pat-token",
    pull_latest: bool = False,
    source_config: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Run a dlt pipeline for the given source.

    Args:
        source_name: Name of the source config (e.g., "data_swim_notion")
        vault_url: Azure Key Vault URL
        base_path: Base path where vd-dlt repos are located
        destination_path: Path for dlt destination (from mounted lakehouse)
        github_pat_secret: Name of GitHub PAT secret in Key Vault
        pull_latest: Whether to git pull before running
        source_config: Optional source config dict passed directly (skips file read)

    Returns:
        Dict with pipeline results
    """
    import dlt

    # Paths
    client_config_path = f"{base_path}/vd-dlt-client-side-config"
    connectors_path = f"{base_path}/vd-dlt-connectors"
    core_path = f"{connectors_path}/core"

    print(f"Source: {source_name}")
    print(f"Base path: {base_path}")

    # Pull latest if requested (only makes sense if not using passed config)
    if pull_latest and not source_config:
        print("\nPulling latest from git...")
        # Get PAT from Key Vault
        import notebookutils
        pat = notebookutils.credentials.getSecret(vault_url, github_pat_secret)
        git_pull(client_config_path, pat)
        git_pull(connectors_path, pat)

    # Add paths for imports
    if core_path not in sys.path:
        sys.path.insert(0, core_path)
    if connectors_path not in sys.path:
        sys.path.insert(0, connectors_path)

    # Load source config - use passed config or read from file
    if source_config:
        print(f"\nUsing provided source config directly")
    else:
        source_config_path = f"{client_config_path}/sources/{source_name}.yml"
        print(f"\nLoading source config: {source_config_path}")
        with open(source_config_path, 'r') as f:
            source_config = yaml.safe_load(f)

    connector_name = source_config.get('connector', 'notion')
    target = source_config.get('target', {})

    print(f"Connector: {connector_name}")
    print(f"Target: {target.get('workspace')}/{target.get('lakehouse')}/{target.get('schema_name')}")

    # Resolve credentials from source config 'secrets' dict
    # Values are Key Vault secret names — resolve each from KV
    secrets = source_config.get('secrets', {})
    credentials = {}

    print(f"\nResolving secrets from Key Vault")
    import notebookutils
    for key, secret_name in secrets.items():
        if isinstance(secret_name, str) and secret_name:
            # Strip legacy ${...} wrapper if present
            if secret_name.startswith("${") and secret_name.endswith("}"):
                secret_name = secret_name[2:-1]
            print(f"  Resolving: {key}")
            credentials[key] = notebookutils.credentials.getSecret(vault_url, secret_name)
        else:
            credentials[key] = secret_name

    # Load connector defaults
    defaults_path = f"{connectors_path}/{connector_name}/defaults.yml"
    connector_defaults = {}
    if os.path.exists(defaults_path):
        with open(defaults_path, 'r') as f:
            connector_defaults = yaml.safe_load(f)
        print(f"\nLoaded connector defaults")

    # Import connector module
    source_module_path = f"{connectors_path}/{connector_name}/source.py"
    spec = importlib.util.spec_from_file_location(f"{connector_name}.source", source_module_path)
    connector_module = importlib.util.module_from_spec(spec)
    sys.modules[f"{connector_name}.source"] = connector_module
    spec.loader.exec_module(connector_module)
    print(f"Loaded connector: {connector_name}")

    # Setup destination - use provided path (from mounted lakehouse or default)
    print(f"Using destination: {destination_path}")
    print(f"  Target config: {target.get('workspace')}/{target.get('lakehouse')}/{target.get('schema_name')}")

    os.environ["DESTINATION__FILESYSTEM__BUCKET_URL"] = destination_path
    os.environ["DESTINATION__FILESYSTEM__DELTALAKE_STORAGE_OPTIONS"] = '{"allow_unsafe_rename": "true"}'

    # Create source
    print("\nCreating dlt source...")
    source = connector_module.create_source(
        credentials=credentials,
        source_config=source_config,
        connector_defaults=connector_defaults
    )

    # Run pipeline
    pipeline_name = source_config.get('pipeline_name', f"{source_name}_pipeline")
    dataset_name = target.get('schema_name', source_name)

    write_disposition = source_config.get('default_sync', {}).get('write_disposition', 'merge')

    print(f"\nRunning pipeline...")
    print(f"  Pipeline: {pipeline_name}")
    print(f"  Dataset: {dataset_name}")
    print(f"  Write disposition: {write_disposition}")

    pipeline = dlt.pipeline(
        pipeline_name=pipeline_name,
        destination="filesystem",
        dataset_name=dataset_name,
    )

    load_info = pipeline.run(
        source,
        write_disposition=write_disposition,
        table_format="delta"
    )

    # Get results
    results = {
        "status": "success",
        "pipeline_name": pipeline_name,
        "dataset_name": dataset_name,
        "destination": f"{destination_path}/{dataset_name}/",
        "load_info": str(load_info),
    }

    # Get row counts
    try:
        row_counts = pipeline.last_trace.last_normalize_info.row_counts
        total = 0
        table_counts = {}
        for table, rows in row_counts.items():
            if not table.startswith("_dlt"):
                table_counts[table] = rows
                total += rows
        results["row_counts"] = table_counts
        results["total_rows"] = total
    except Exception as e:
        results["row_counts_error"] = str(e)

    # Print summary
    print("\n" + "=" * 60)
    print("COMPLETED")
    print("=" * 60)
    print(load_info)

    if "row_counts" in results:
        print("\nRows loaded:")
        for table, rows in results["row_counts"].items():
            print(f"  {table}: {rows}")
        print(f"\nTotal: {results['total_rows']} rows")

    print(f"\nDestination: {results['destination']}")

    return results


def run_from_notebook(
    source_name: str,
    vault_url: str,
    base_path: str = "/tmp/vd-dlt",
    github_pat_secret: str = "my-github-pat-token",
    pull_latest: bool = False,
    source_config_json: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Complete pipeline run from Fabric notebook.

    Does everything after repos are cloned:
    - Loads source config (from file or passed JSON)
    - Mounts target lakehouse
    - Runs the pipeline

    Args:
        source_name: Name of the source config (e.g., "data_swim_notion")
        vault_url: Azure Key Vault URL
        base_path: Base path where vd-dlt repos are cloned
        github_pat_secret: Name of GitHub PAT secret in Key Vault
        pull_latest: Whether to git pull before running
        source_config_json: Optional JSON string of source config (skips file read)

    Returns:
        Dict with pipeline results
    """
    import json

    # Load source config - from JSON param or file
    if source_config_json:
        print("Using source config from parameter")
        source_config = json.loads(source_config_json)
    else:
        source_config_path = f"{base_path}/vd-dlt-client-side-config/sources/{source_name}.yml"
        with open(source_config_path, 'r') as f:
            source_config = yaml.safe_load(f)

    target = source_config.get('target', {})
    lakehouse_name = target.get('lakehouse')
    lakehouse_id = target.get('lakehouse_id')
    workspace_id = target.get('workspace_id')

    print(f"Target lakehouse: {lakehouse_name}")
    print(f"  Lakehouse ID: {lakehouse_id}")
    print(f"  Workspace ID: {workspace_id}")

    # Mount lakehouse
    destination_path = mount_lakehouse(workspace_id, lakehouse_id)

    # Run pipeline
    return run_pipeline(
        source_name=source_name,
        vault_url=vault_url,
        base_path=base_path,
        destination_path=destination_path,
        github_pat_secret=github_pat_secret,
        pull_latest=pull_latest,
        source_config=source_config if source_config_json else None,
    )
