import pytest

from soigia.base_pipeline import BasePipeline, Pipeline, PipelineStatus


def test_base_pipeline_decorator_and_hooks():
    pipeline = BasePipeline(name="demo")
    events = []

    @pipeline.before_run
    def before(ctx):
        events.append(("before_run", ctx.run_id))

    @pipeline.after_run
    def after(ctx):
        events.append(("after_run", ctx.run_id))

    @pipeline.step(name="double", retries=1)
    def double(data, ctx):
        events.append(("step", ctx.current_step, ctx.current_attempt))
        return data * 2

    outcome = pipeline.run(10)

    assert outcome.status == PipelineStatus.SUCCESS
    assert outcome.data == 20
    assert events[0][0] == "before_run"
    assert events[-1][0] == "after_run"
    assert events[1] == ("step", "double", 1)


def test_base_pipeline_condition_and_skip():
    pipeline = BasePipeline(name="demo")

    @pipeline.step(name="skip_me", condition=lambda data, ctx: False)
    def skipped(data, ctx):  # pragma: no cover - this should not run
        return data + 1

    outcome = pipeline.run(3)

    assert outcome.status == PipelineStatus.SUCCESS
    assert outcome.data == 3
    assert pipeline.steps[0].name == "skip_me"


def test_pipeline_subclass_auto_registers_decorated_methods():
    class DemoPipeline(Pipeline):
        stages = ["double", "add_three"]

        def load_data(self):
            return 4

        def double(self):
            self.data_df = self.data_df * 2
            return self.data_df

        def add_three(self):
            self.data_df = self.data_df + 3
            return self.data_df

    pipeline = DemoPipeline(name="demo")

    assert [step.name for step in pipeline.steps] == ["double", "add_three"]
    outcome = pipeline.run()
    assert outcome.status == PipelineStatus.SUCCESS
    assert outcome.data == 11


def test_pipeline_exposes_shared_data_df_and_model():
    class DemoPipeline(Pipeline):
        stages = ["capture", "finalize"]

        def load_data(self):
            return 4

        def capture(self):
            assert self.data_df == 4
            self.model.seen = getattr(self.model, "seen", 0) + 1
            self.data_df = self.data_df + 1
            return self.data_df

        def finalize(self):
            assert self.model.seen == 1
            assert self.data_df == 5
            self.data_df = self.data_df * 2
            return self.data_df

    pipeline = DemoPipeline(name="demo")
    outcome = pipeline.run()

    assert outcome.status == PipelineStatus.SUCCESS
    assert outcome.data == 10
    assert pipeline.data_df == 10
    assert pipeline.model.seen == 1


def test_pipeline_helper_methods():
    import pandas as pd
    from pathlib import Path

    class DemoPipeline(Pipeline):
        stages = ["prepare", "finalize"]

        def load_data(self):
            import pandas as pd

            return pd.DataFrame({"name": ["alice"], "amount": [3]})

        def prepare(self):
            assert self.has_columns("name", "amount")
            self.snapshot_df("before")
            self.put("note", "ready")
            self.data_df = self.data_df.assign(amount2=self.data_df["amount"] * 2)
            return self.data_df

        def finalize(self):
            assert self.get("note") == "ready"
            self.restore_df("before")
            assert "amount2" not in self.data_df.columns
            self.data_df = self.data_df.assign(done=True)
            return self.data_df

    pipeline = DemoPipeline(name="demo")
    result = pipeline.run()

    assert result.status == PipelineStatus.SUCCESS
    assert result.rows == 1
    assert result.csv_path is not None
    assert result.db_path is not None
    assert result.parquet_path is not None
    assert result.summary_path is not None
    assert result.duration_seconds is not None
    assert pipeline.df.equals(pd.DataFrame({"name": ["alice"], "amount": [3], "done": [True]}))
    dump_dir = Path.cwd() / "data" / "DemoPipeline"
    assert (dump_dir / "01_prepare.csv").exists()
    assert (dump_dir / "02_finalize.csv").exists()
    summary_path = dump_dir / "pipeline_summary.md"
    summary_text = summary_path.read_text(encoding="utf-8", errors="replace")
    assert "# Pipeline Run Summary" in summary_text
    assert "| # | Step | Status | Rows | Added Columns | Removed Columns |" in summary_text
    assert "| 01 | prepare | success |" in summary_text
    assert "| 02 | finalize | success |" in summary_text
    assert "## Model" in summary_text


def test_pipeline_before_and_after_methods_are_called():
    class DemoPipeline(Pipeline):
        stages = ["step_one"]

        def load_data(self):
            return 0

        def before_run(self):
            self.model.events = ["before"]

        def after_run(self):
            self.model.events.append("after")

        def step_one(self):
            self.model.events.append("step")
            self.data_df = 1
            return self.data_df

    pipeline = DemoPipeline(name="demo")
    outcome = pipeline.run()

    assert outcome.status == PipelineStatus.SUCCESS
    assert pipeline.model.events == ["before", "step", "after"]


def test_pipeline_module_entrypoint_import():
    from soigia.pipeline import BasePipeline as PublicBasePipeline

    assert PublicBasePipeline is Pipeline


def test_pipeline_loaders_support_data_csv_and_yaml():
    import pandas as pd
    import yaml
    from pathlib import Path

    class LoaderPipeline(Pipeline):
        stages = ["normalize"]

        def load_data(self):
            return [{"name": " alice ", "amount": "12.5"}, {"name": "bob", "amount": "7"}]

        def normalize(self):
            self.logging("loaders-test-marker")
            df = pd.DataFrame(self.require_df())
            df["name"] = df["name"].astype(str).str.strip()
            df["amount"] = pd.to_numeric(df["amount"], errors="coerce")
            self.data_df = df
            return self.data_df

    data_pipeline = LoaderPipeline(name="data")
    data_result = data_pipeline.run()
    assert data_result.success
    assert data_pipeline.df.shape == (2, 2)
    dump_dir = Path.cwd() / "data" / "LoaderPipeline"
    assert (dump_dir / "01_normalize.csv").exists()
    log_path = Path.cwd() / "logs" / "LoaderPipeline.log"
    assert log_path.exists()
    assert "loaders-test-marker" in log_path.read_text(encoding="utf-8", errors="replace")
    summary_text = (dump_dir / "pipeline_summary.md").read_text(encoding="utf-8", errors="replace")
    assert "# Pipeline Run Summary" in summary_text
    assert "| 01 | normalize | success |" in summary_text

    artifact_root = Path.cwd() / ".test_artifacts"
    artifact_root.mkdir(exist_ok=True)

    csv_path = artifact_root / "sales-loader.csv"
    pd.DataFrame([{"name": "alice", "amount": 12.5}, {"name": "bob", "amount": 7.0}]).to_csv(
        csv_path, index=False
    )

    class CsvPipeline(Pipeline):
        stages = ["normalize"]

        def load_data(self):
            return self.load_csv(csv_path)

        def normalize(self):
            self.data_df["name"] = self.data_df["name"].str.upper()
            return self.data_df

    csv_pipeline = CsvPipeline(name="csv")
    csv_result = csv_pipeline.run()
    assert csv_result.success
    assert list(csv_pipeline.df["name"]) == ["ALICE", "BOB"]
    assert (Path.cwd() / "data" / "CsvPipeline" / "01_normalize.csv").exists()

    yaml_path = artifact_root / "sales-loader.yaml"
    yaml_payload = [{"name": "carol", "amount": 9.5}, {"name": "dave", "amount": 3.25}]
    yaml_path.write_text(
        yaml.safe_dump(yaml_payload, sort_keys=False, allow_unicode=True),
        encoding="utf-8",
    )

    class YamlPipeline(Pipeline):
        stages = ["normalize"]

        def load_data(self):
            payload = self.load_yaml(yaml_path)
            self.data_df = pd.DataFrame(payload)
            return self.data_df

        def normalize(self):
            self.data_df["amount"] = pd.to_numeric(self.data_df["amount"], errors="coerce")
            return self.data_df

    yaml_pipeline = YamlPipeline(name="yaml")
    yaml_result = yaml_pipeline.run()
    assert yaml_result.success
    assert yaml_pipeline.df.shape == (2, 2)
    assert (Path.cwd() / "data" / "YamlPipeline" / "01_normalize.csv").exists()


def test_pipeline_pipelines_alias_still_works():
    class AliasPipeline(Pipeline):
        pipelines = ["double"]

        def load_data(self):
            return 2

        def double(self):
            self.data_df = self.data_df * 2
            return self.data_df

    pipeline = AliasPipeline(name="alias")
    outcome = pipeline.run()
    assert outcome.success
    assert outcome.data == 4


def test_pipeline_protected_hooks_cannot_be_overridden():
    with pytest.raises(TypeError):

        class BadPreprocessPipeline(Pipeline):
            def preprocess_data(self):  # pragma: no cover - defined only to trigger the guard
                return self.data_df

    with pytest.raises(TypeError):

        class BadPostprocessPipeline(Pipeline):
            def postprocess_data(self):  # pragma: no cover - defined only to trigger the guard
                return self.data_df

    with pytest.raises(TypeError):

        class BadSavePipeline(Pipeline):
            def save_data_to_db(self):  # pragma: no cover - defined only to trigger the guard
                return self.data_df

    with pytest.raises(TypeError):

        class BadSaveOutputsPipeline(Pipeline):
            def save_outputs(self):  # pragma: no cover - defined only to trigger the guard
                return self.data_df


def test_pipeline_defaults_to_empty_dataframe_when_load_data_not_overridden():
    class EmptyPipeline(Pipeline):
        stages = []

    pipeline = EmptyPipeline(name="empty")
    result = pipeline.run()

    assert result.success
    assert result.rows == 0
    assert hasattr(pipeline.data_df, "empty")
    assert pipeline.data_df.empty
    assert result.csv_path is not None
    assert result.db_path is None
    assert result.parquet_path is None


def test_pipeline_uses_load_data_return_value():
    class ReturnPipeline(Pipeline):
        stages = ["double"]

        def load_data(self):
            return 3

        def double(self):
            self.data_df = self.data_df * 2
            return self.data_df

    pipeline = ReturnPipeline(name="return")
    result = pipeline.run()

    assert result.success
    assert pipeline.data_df == 6


def test_pipeline_loads_yaml_config_into_namespace():
    import yaml
    from pathlib import Path

    config_dir = Path(__file__).resolve().parent
    config_path = config_dir / "config.yaml"
    config_path.write_text(
        yaml.safe_dump({"factor": 4, "label": "demo", "nested": {"enabled": True}}, sort_keys=False),
        encoding="utf-8",
    )

    class YamlConfigPipeline(Pipeline):
        stages = ["scale"]

        def load_data(self):
            return 3

        def scale(self):
            self.data_df = self.data_df * self.config.factor
            self.model.label = self.config.label
            self.model.enabled = self.config.nested["enabled"]
            return self.data_df

    pipeline = YamlConfigPipeline(name="yaml-config")
    result = pipeline.run()

    assert result.success
    assert pipeline.data_df == 12
    assert pipeline.config.factor == 4
    assert pipeline.config.label == "demo"
    assert pipeline.config.nested["enabled"] is True
    assert pipeline.config.get("factor") == 4
    config_path.unlink(missing_ok=True)


def test_pipeline_restores_snapshot_when_stage_fails():
    import pandas as pd

    class FailingPipeline(Pipeline):
        stages = ["prepare", "explode"]

        def load_data(self):
            return pd.DataFrame({"value": [1, 2]})

        def prepare(self):
            self.data_df = self.data_df.assign(prepared=True)
            return self.data_df

        def explode(self):
            self.data_df = self.data_df.assign(bad=True)
            raise RuntimeError("boom")

    pipeline = FailingPipeline(name="failing")
    result = pipeline.run()

    assert result.failed
    assert list(pipeline.data_df.columns) == ["value", "prepared"]
    assert "bad" not in pipeline.data_df.columns


def test_pipeline_executes_registered_rollback_actions_on_failure():
    class SideEffectPipeline(Pipeline):
        stages = ["prepare", "explode"]

        def load_data(self):
            return 1

        def prepare(self):
            self.model.events = []
            self.model.events.append("prepare")
            self.on_rollback(lambda: self.model.events.append("undo"))
            self.data_df = 2
            return self.data_df

        def explode(self):
            self.model.events.append("explode")
            raise RuntimeError("boom")

    pipeline = SideEffectPipeline(name="side-effect")
    result = pipeline.run()

    assert result.failed
    assert pipeline.model.events == ["prepare", "explode", "undo"]
