"""Stateless orchestrator: print the recommended diff-diff workflow with
the caller's column names wired in.

This module exists to give LLM agents a single, recognizable entrypoint
that names the rest of the agent-facing workflow (`profile_panel`,
`get_llm_guide`, `practitioner_next_steps`, `BusinessReport`). The
function does not fit, inspect, or recommend — it templates a copy-
pasteable script.
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple

# Pattern → df-callable estimator class names. Flat union below is the
# `fit_candidates` field of the returned dict; each name must remain a
# valid `hasattr(diff_diff, name)` (locked by the contract test in
# tests/test_agent_discoverability.py and tests/test_agent_workflow.py).
# Patterns intentionally exclude post-fit and pre-fit diagnostics
# (PreTrendsPower takes pre-treatment coefficients, HonestDiD takes a
# fitted results object); those are mentioned separately in the
# templated Step 4 of the script.
_WORKFLOW_PATTERNS: Tuple[Tuple[str, Tuple[str, ...]], ...] = (
    (
        "Staggered adoption + binary treatment + has_never_treated control",
        ("CallawaySantAnna", "SunAbraham", "ImputationDiD"),
    ),
    (
        "Continuous treatment dose (non-binary numeric intensity)",
        ("ContinuousDiD",),
    ),
    (
        "Heterogeneous adoption intensity across treated units",
        ("HeterogeneousAdoptionDiD",),
    ),
    (
        "Simple 2x2 DiD (binary treatment, two periods, no staggering)",
        ("DifferenceInDifferences",),
    ),
)


def _safe_kwarg(name: str, value: Optional[str]) -> Optional[str]:
    """Render ``name=<python-literal>`` using repr() for source-safety.

    Column labels containing quotes, backslashes, or other special
    characters must not break the emitted "copy-pasteable" script.
    Python's built-in ``repr()`` produces a valid string literal for
    any str input (including embedded quotes / backslashes /
    newlines), so ``f"{name}={value!r}"`` is injection-safe by
    construction. ``None`` returns ``None`` so the caller can drop
    the kwarg.
    """
    if value is None:
        return None
    return f"{name}={value!r}"


def _join_kwargs(**kwargs: Optional[str]) -> str:
    parts = [_safe_kwarg(k, v) for k, v in kwargs.items()]
    return ", ".join(p for p in parts if p is not None)


def agent_workflow(
    df: Any,
    *,
    unit: str,
    time: str,
    treatment: str,
    outcome: str,
    first_treat: Optional[str] = None,
    df_name: str = "df",
    verbose: bool = True,
) -> Dict[str, Any]:
    """Print the recommended diff-diff workflow with your column names wired in.

    Stateless orchestrator. Calls nothing internally. Returns a dict;
    optionally prints a copy-pasteable script (``verbose=True``, the
    default). ``df`` is not inspected — column names are templated
    verbatim into the output.

    Parameters
    ----------
    df : pandas.DataFrame
        Long-format panel data. Not inspected; included so the agent
        can pass the same handle along to the next call.
    unit : str
        Column identifying the cross-sectional unit.
    time : str
        Column identifying the time period.
    treatment : str
        Column holding the treatment indicator or dose.
    outcome : str
        Column holding the outcome variable.
    first_treat : str, optional
        Column with each unit's first-treatment period (or NaN for
        never-treated controls). When supplied, the templated Step 3
        switches from a ``DifferenceInDifferences.fit(treatment=...)``
        example to a ``CallawaySantAnna().fit(first_treat=...)``
        example, matching the actual fit signatures (passing
        ``treatment=`` to CallawaySantAnna's ``.fit()`` would raise
        TypeError).
    df_name : str, default ``"df"``
        Identifier under which the caller's dataframe is bound in
        their namespace. Templated verbatim into the emitted script
        as the first positional argument of every call
        (``profile_panel({df_name}, ...)``,
        ``<Estimator>().fit({df_name}, ...)``) so the script is
        directly executable when the caller's local variable matches.
        If the caller has ``panel = pd.read_parquet(...)``, passing
        ``df_name="panel"`` produces a script that references
        ``panel`` instead of ``df``. Must be a valid Python identifier
        (not enforced; non-identifier values produce a script that
        won't parse).
    verbose : bool, default True
        If True, print the script to stdout. The dict is always
        returned regardless.

    Returns
    -------
    dict
        Keys:

        - ``"profile_call"`` (str): call signature for
          :func:`diff_diff.profile_panel`.
        - ``"guide_call"`` (str): call signature for
          :func:`diff_diff.get_llm_guide`.
        - ``"fit_candidates"`` (list of str): flat union of estimator /
          diagnostic class names referenced in the workflow patterns.
          Every name resolves on the top-level ``diff_diff`` namespace.
        - ``"validation_calls"`` (list of str): call signatures for the
          post-fit validation step.
        - ``"reporting_call"`` (str): call signature for
          :class:`diff_diff.BusinessReport`.
        - ``"script"`` (str): printable multi-line workflow.

    Examples
    --------
    >>> import pandas as pd
    >>> import diff_diff
    >>> df = pd.DataFrame({
    ...     "firm_id": [1, 1, 2, 2],
    ...     "year": [0, 1, 0, 1],
    ...     "treated": [0, 0, 1, 1],
    ...     "logwage": [0.1, 0.2, 0.1, 0.9],
    ... })
    >>> out = diff_diff.agent_workflow(df, unit="firm_id", time="year",
    ...                                treatment="treated", outcome="logwage",
    ...                                verbose=False)
    >>> "profile_panel" in out["script"]
    True
    """
    del df  # intentionally unused: orchestrator templates from column names only

    profile_call = (
        f"diff_diff.profile_panel({df_name}, "
        f"{_join_kwargs(unit=unit, time=time, treatment=treatment, outcome=outcome)})"
    )
    guide_call = 'diff_diff.get_llm_guide("autonomous")'

    # Step 3 example: branch on first_treat presence.
    # - With first_treat: a staggered structure is strongly implied, BUT
    #   `first_treat` does not by itself identify which estimator to use:
    #   CallawaySantAnna (binary staggered), ContinuousDiD (continuous-
    #   dose with first_treat), and HeterogeneousAdoptionDiD event-study
    #   (heterogeneous intensity with first_treat_col) all accept it.
    #   Show CallawaySantAnna as the binary-staggered canonical example
    #   and list the alternatives for continuous / heterogeneous designs
    #   so an agent isn't steered to the wrong estimator.
    # - Without first_treat: the orchestrator does not inspect df, so it
    #   CANNOT infer whether the panel is 2x2 binary vs continuous-dose
    #   vs heterogeneous-adoption. Show a DifferenceInDifferences call
    #   as the "simple 2x2" example and label it explicitly conditional
    #   on that shape.
    if first_treat is not None:
        fit_example_kwargs = _join_kwargs(
            outcome=outcome, unit=unit, time=time, first_treat=first_treat
        )
        fit_example_call = f"diff_diff.CallawaySantAnna().fit({df_name}, {fit_example_kwargs})"
        step3_label_lines = [
            "Step 3 - Fit. Your data has `first_treat` -> staggered structure.",
            "`first_treat` alone does NOT identify a single estimator; pick by",
            "treatment shape:",
            "  - Binary staggered  : CallawaySantAnna (shown) / SunAbraham / ImputationDiD",
            "  - Continuous dose   : ContinuousDiD (also takes first_treat=)",
            "  - Heterogeneous adoption intensity:",
            "                        HeterogeneousAdoptionDiD (event study,",
            "                        takes first_treat_col=, NOT first_treat=)",
        ]
    else:
        fit_example_kwargs = _join_kwargs(
            outcome=outcome, unit=unit, time=time, treatment=treatment
        )
        fit_example_call = (
            f"diff_diff.DifferenceInDifferences().fit({df_name}, {fit_example_kwargs})"
        )
        step3_label_lines = [
            "Step 3 - Fit. Pick a candidate from Step 2's patterns based on your",
            "treatment/time shape. The example below shows the simple 2x2 case",
            "(binary treatment + binary time); substitute ContinuousDiD /",
            "HeterogeneousAdoptionDiD / etc. when your design is not 2x2",
            "(DifferenceInDifferences.fit() validates and rejects non-binary",
            "treatment or time).",
        ]
    step3_comment_block = "\n".join(f"# {line}" for line in step3_label_lines)

    validation_calls = [
        "diff_diff.practitioner_next_steps(result)",
    ]
    reporting_call = "diff_diff.BusinessReport(result).full_report()"

    fit_candidates: List[str] = []
    pattern_lines: List[str] = []
    for label, names in _WORKFLOW_PATTERNS:
        pattern_lines.append(f"#   - {label}")
        pattern_lines.append(f"#       candidates: {', '.join(names)}")
        for n in names:
            if n not in fit_candidates:
                fit_candidates.append(n)
    pattern_block = "\n".join(pattern_lines)

    diagnostics_block = (
        "# Parallel-trends sensitivity / power (take a fitted result or\n"
        "# pre-trend coefficients, NOT df+columns):\n"
        "#     diff_diff.PreTrendsPower / diff_diff.HonestDiD"
    )

    # Templated output is a valid Python script: every prose line is a
    # `#` comment, every code line stands at column 0 and runs as-is.
    # Step 5 wraps full_report() in print() so end-to-end execution
    # actually produces the stakeholder narrative.
    script = f"""# diff_diff workflow for your data
# =================================
#
# Step 1 - Describe the panel:
profile = {profile_call}
print(profile)

# Step 2 - Choose an estimator. Consult the routing matrix:
print({guide_call})

# Routing patterns (df-callable estimators):
{pattern_block}
#
{diagnostics_block}

{step3_comment_block}
result = {fit_example_call}

# Step 4 - Validate:
{validation_calls[0]}

# Step 5 - Report:
print({reporting_call})

# Full reference: diff_diff.get_llm_guide("full")
# Practitioner recipe: diff_diff.get_llm_guide("practitioner")
"""

    if verbose:
        print(script)

    return {
        "profile_call": profile_call,
        "guide_call": guide_call,
        "fit_candidates": fit_candidates,
        "validation_calls": validation_calls,
        "reporting_call": reporting_call,
        "script": script,
    }
