from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal, Union, Dict, Optional, Any
from asyncio import gather
from asyncer import syncify
from pyspark.sql import DataFrame

from transforms.api.transform_df import Transform, TransformOutput
from transforms.runner.data_sink.base import DataSink
from transforms.runner.data_source.base import DataSource

from .exec_check import execute_check
import logging
logger = logging.getLogger(__name__)

@dataclass
class TransformRunner:
    sourcer: DataSource
    sink: DataSink
    fallback_branches: list[str] = field(default_factory=list)
    output_dir: Path = Path.home() / ".fndry_duck" / "output"
    secrets_config_location: Path = Path.home() / ".fndry_duck" / "secrets"
    
    def __post_init__(self):
        self.output_dir.mkdir(parents=True, exist_ok=True)

    async def download_datasets(self, transform: Transform, omit_checks: bool, dry_run: bool) -> Dict[str, Union[DataFrame, TransformOutput]]:
        sources: Dict[str, Union[DataFrame, TransformOutput]] = {}
        futures_list = [(argname, self.sourcer.download_for_branches(input.path_or_rid, branches=[
                b for b in ([input.branch] + self.fallback_branches) if b is not None
            ])) for argname, input in transform.inputs.items()]
        futures = [fut for (_, fut) in futures_list]
        names = [name for (name, _) in futures_list]
        dfs: list[DataFrame] = await gather(*futures)
        for argname, res_df in zip(names, dfs):
            sources[argname] = res_df
        return sources
            
    def exec_transform(self, transform: Transform, omit_checks: bool, dry_run: bool, sources: Optional[Dict[str, Union[DataFrame, TransformOutput]]] = None) -> None:
        if sources is None:
            sources = syncify(self.download_datasets, raise_sync_error=False)(transform, omit_checks, dry_run)
            
        if transform.multi_outputs is not None:
            impl_multi_outputs = {}
            for argname, output in transform.outputs.items():
                def on_dataframe_req(mode: Literal["current", "previous"]) -> DataFrame:
                    if mode == "current":
                        return syncify(self.sourcer.download_for_branches)(
                            output.path_or_rid, branches=self.fallback_branches
                        )
                    elif (transform.incremental_opts is not None) and (not dry_run):
                        return syncify(self.sourcer.download_latest_incremental_transaction)(
                            dataset_path_or_rid=output.path_or_rid,
                            branches=self.fallback_branches,
                            semantic_version=transform.incremental_opts.semantic_version
                        )
                    raise ValueError(f"Invalid mode {mode} or missing incremental options")

                # Create a closure to capture the current output variable
                def create_write_fn(output_path: str):
                    def on_dataframe_write(df: DataFrame, mode: Literal["append", "replace"]) -> None:
                        if mode == "append":
                            raise NotImplementedError()
                        else:
                            if not omit_checks:
                                for check in output.checks: 
                                    execute_check(df, check)
                            if (transform.incremental_opts is not None) and (not dry_run):
                                self.sink.save_incremental_transaction(
                                    df,
                                    output_path,
                                    transform.incremental_opts.semantic_version
                                )
                            else:
                                logger.info(f"Saving output to {output_path}")
                                self.sink.save_transaction(
                                    df=df,
                                    dataset_path_or_rid=output_path
                                )
                    return on_dataframe_write

                output_df_impl = TransformOutput(
                    on_dataframe_req=on_dataframe_req,
                    on_dataframe_write=create_write_fn(output.path_or_rid)
                )
                impl_multi_outputs[argname] = output_df_impl
                sources[argname] = output_df_impl
            
            # Set the multi_outputs dictionary with our implementations
            transform.multi_outputs = impl_multi_outputs
        
        logger.info("Starting transform")
        res = transform.transform(**sources)
        logger.info("Finished transform")
        
        if transform.multi_outputs is None:
            # Accept any DataFrame-like object (including engine-specific SQLFrame DataFrames)
            if not hasattr(res, "limit") or not hasattr(res, "collect"):
                raise ValueError("Transform without multi_outputs must return a DataFrame-like object")
            
            if not omit_checks:
                for check in transform.outputs["output"].checks:
                    logger.info(f"Running check {check.description}")
                    execute_check(res if not dry_run else res.limit(1), check)
                    logger.info(f"Check {check.description} finished")
            if not dry_run:
                logger.info(f"Saving transaction to {transform.outputs['output'].path_or_rid}")
                self.sink.save_transaction(df=res, dataset_path_or_rid=transform.outputs['output'].path_or_rid)
                logger.info("Transaction saved successfully")
            else:
                res.limit(1).collect()
        else:
            if transform.incremental_opts is not None:
                logger.info("Finished transform successfully")
                raise NotImplementedError("Not yet implemented saving for incremental dataset")
            else:
                logger.info("Finished transform successfully")
                
