scider/workflows/full_workflow.py

"""
Full SciDER Workflow

Complete workflow that chains DataAgent and ExperimentAgent from scratch.
This workflow takes raw data, analyzes it, generates experiment code,
executes it, and produces final metrics.

For partial workflows (e.g., starting from existing data analysis), see:
- data_workflow.py: Only runs DataAgent
- experiment_workflow.py: Only runs ExperimentAgent
"""

import shutil
from pathlib import Path
from typing import Literal

from loguru import logger
from pydantic import BaseModel, PrivateAttr, model_validator

from scider.core.code_env import WorkspaceInitConfig
from scider.core.constant import override_user_approval
from scider.workflows.data_workflow import DataWorkflow
from scider.workflows.hypo_data_workflow import HypoDataWorkflow
from scider.workflows.experiment_workflow import ExperimentWorkflow
from scider.workflows.paper_bootstrap import build_experimental_log, build_sparse_idea_from_query
from scider.workflows.utils import get_separator
from scider.workflows.writing_workflow import WritingWorkflow


class FullWorkflow(BaseModel):
    """
    Full SciDER Workflow - chains DataAgent and ExperimentAgent from scratch.

    This workflow executes:
    1. DataWorkflow - Analyzes input data, produces data_analysis.md
    2. ExperimentWorkflow - Generates code, executes experiments, produces metrics

    Internally uses DataWorkflow and ExperimentWorkflow for better modularity.

    Usage:
        workflow = FullWorkflow(
            data_path="data/data.csv",
            workspace_path="workspace",
            user_query="Train an SVR model",
        )
        workflow.run()
        print(workflow.final_summary)
    """

    # ==================== INPUT ====================
    data_path: Path | None = None        # real data; mutually exclusive with feature_desc
    feature_desc: str | None = None      # synthetic data description; mutually exclusive with data_path
    num_rows: int = 1000                 # forwarded to HypoDataWorkflow
    workspace_path: Path
    user_query: str
    repo_source: str | None = None
    max_revisions: int = 5
    data_agent_recursion_limit: int = 100
    experiment_agent_recursion_limit: int = 100
    data_desc: str | None = None  # Optional additional description of the data

    # Paper writing phase (off by default)
    run_paper_writing: bool = False
    paper_workspace_path: Path | None = None
    paper_template_dir_path: Path | None = None
    paper_template_tex_path: Path | None = None
    paper_conference_guidelines_path: Path | None = None
    paper_agent_recursion_limit: int = 150

    # None → LocalEnv's default (uv-managed, auto `uv init`). Forwarded to
    # the nested DataWorkflow / ExperimentWorkflow / WritingWorkflow.
    workspace_init_config: WorkspaceInitConfig | None = None

    # ==================== INTERNAL STATE ====================
    current_phase: Literal[
        "init", "data_analysis", "experiment", "paper_writing", "complete", "failed"
    ] = "init"
    data_summary: str = ""
    data_agent_history: list = []

    # Paper subagent results (from DataWorkflow)
    papers: list[dict] = []
    datasets: list[dict] = []
    metrics: list[dict] = []
    paper_search_summary: str | None = None

    # ==================== OUTPUT ====================
    final_status: Literal["success", "failed", "max_revisions_reached"] | None = None
    final_summary: str = ""
    execution_results: list = []
    error_message: str | None = None

    # Paper writing outputs (populated only if run_paper_writing=True)
    paper_final_tex_path: Path | None = None
    paper_final_pdf_path: Path | None = None
    paper_writing_summary: str = ""

    # Internal: sub-workflows
    _data_workflow: DataWorkflow | HypoDataWorkflow | None = PrivateAttr(default=None)
    _experiment_workflow: ExperimentWorkflow | None = PrivateAttr(default=None)
    _writing_workflow: "WritingWorkflow | None" = PrivateAttr(default=None)

    @model_validator(mode="after")
    def _validate_data_mode(self) -> "FullWorkflow":
        has_path = self.data_path is not None
        has_desc = bool(self.feature_desc and self.feature_desc.strip())
        if has_path and has_desc:
            raise ValueError(
                "Provide either data_path (real data) or feature_desc (synthetic), not both."
            )
        if not has_path and not has_desc:
            raise ValueError("FullWorkflow requires either data_path or feature_desc.")
        return self

    def run(self) -> "FullWorkflow":
        """
        Run the complete workflow: DataWorkflow -> ExperimentWorkflow.

        Returns:
            self (for chaining)
        """

        logger.info(get_separator())
        logger.info("Starting Full SciDER Workflow")
        logger.info(get_separator())

        # Step 1: Run DataWorkflow
        if not self._run_data_phase():
            self._finalize()
            return self

        # Step 2: Run ExperimentWorkflow
        self._run_experiment_phase()

        # Step 3: (Optional) Paper writing
        if self.run_paper_writing and self.current_phase != "failed":
            self._run_paper_writing_phase()

        # Step 4: Finalize
        self._finalize()

        return self

    def _run_data_phase(self) -> bool:
        """
        Run DataWorkflow or HypoDataWorkflow depending on which input was provided.

        Returns:
            True if successful, False if failed
        """
        logger.info("Phase 1: Running data phase")
        self.current_phase = "data_analysis"

        if self.feature_desc:
            self._data_workflow = HypoDataWorkflow(
                feature_desc=self.feature_desc,
                workspace_path=self.workspace_path,
                num_rows=self.num_rows,
                user_query=self.user_query,
                recursion_limit=self.data_agent_recursion_limit,
                extra_data_desc=self.data_desc,
                workspace_init_config=self.workspace_init_config,
            )
        else:
            self._data_workflow = DataWorkflow(
                data_path=self.data_path,
                workspace_path=self.workspace_path,
                recursion_limit=self.data_agent_recursion_limit,
                data_desc=self.data_desc,
                workspace_init_config=self.workspace_init_config,
            )

        try:
            self._data_workflow.run()

            if self._data_workflow.final_status == "success":
                self.data_summary = self._data_workflow.data_summary
                if isinstance(self._data_workflow, DataWorkflow):
                    self.data_agent_history = self._data_workflow.data_agent_history
                self._data_workflow.save_summary()
                logger.info("Data phase completed successfully")
                return True
            else:
                self.error_message = self._data_workflow.error_message
                self.current_phase = "failed"
                return False

        except Exception as e:
            logger.exception("Data phase failed")
            self.error_message = f"Data phase failed: {e}"
            self.current_phase = "failed"
            return False

    def _run_experiment_phase(self) -> bool:
        """
        Run ExperimentWorkflow to generate and execute experiments.

        Returns:
            True if successful, False if failed
        """
        logger.info("Phase 2: Running ExperimentWorkflow")
        self.current_phase = "experiment"

        self._experiment_workflow = ExperimentWorkflow(
            workspace_path=self.workspace_path,
            user_query=self.user_query,
            data_summary=self.data_summary,
            repo_source=self.repo_source,
            max_revisions=self.max_revisions,
            recursion_limit=self.experiment_agent_recursion_limit,
            workspace_init_config=self.workspace_init_config,
        )

        try:
            self._experiment_workflow.run()

            # Extract results and save summary to workspace
            self.final_status = self._experiment_workflow.final_status
            self.execution_results = self._experiment_workflow.execution_results
            self._experiment_workflow.save_summary()
            self.final_summary = self._compose_summary()
            self.current_phase = "complete"

            logger.info(f"ExperimentWorkflow completed: {self.final_status}")
            return True

        except Exception as e:
            logger.exception("ExperimentWorkflow failed")
            self.error_message = f"ExperimentWorkflow failed: {e}"
            self.current_phase = "failed"
            self.final_status = "failed"
            return False

    def _run_paper_writing_phase(self) -> bool:
        """Run the WritingWorkflow to turn the SciDER outputs into a paper."""
        logger.info("Phase 3: Running WritingWorkflow for paper generation")
        self.current_phase = "paper_writing"

        idea_md = build_sparse_idea_from_query(self.user_query)
        experiment_final_summary = (
            self._experiment_workflow.final_summary if self._experiment_workflow else ""
        )
        experimental_log = build_experimental_log(
            data_summary=self.data_summary,
            experiment_summary=experiment_final_summary,
            user_query=self.user_query,
        )

        return self._execute_writing_workflow(idea_md, experimental_log, experiment_final_summary)

    def _execute_writing_workflow(
        self, idea_md: str, experimental_log: str, experiment_summary: str
    ) -> bool:
        """Shared logic: instantiate, run, and extract results from WritingWorkflow."""
        from scider.workflows.utils import run_paper_writing_phase

        try:
            self._writing_workflow = run_paper_writing_phase(
                workspace_path=self.workspace_path,
                idea_md=idea_md,
                experimental_log=experimental_log,
                user_query=self.user_query,
                data_summary=self.data_summary,
                experiment_summary=experiment_summary,
                paper_workspace_path=self.paper_workspace_path,
                paper_template_dir_path=self.paper_template_dir_path,
                paper_template_tex_path=self.paper_template_tex_path,
                paper_conference_guidelines_path=self.paper_conference_guidelines_path,
                paper_agent_recursion_limit=self.paper_agent_recursion_limit,
                workspace_init_config=self.workspace_init_config,
            )
            if self._writing_workflow.final_status != "success":
                self.error_message = (
                    self._writing_workflow.error_message or "WritingWorkflow failed"
                )
                self.current_phase = "failed"
                return False

            self.paper_final_tex_path = self._writing_workflow.final_tex_path
            self.paper_final_pdf_path = self._writing_workflow.final_pdf_path
            self.paper_writing_summary = self._writing_workflow.final_summary
            logger.info("WritingWorkflow completed successfully")
            return True
        except Exception as e:
            logger.exception("WritingWorkflow failed")
            self.error_message = f"WritingWorkflow failed: {e}"
            self.current_phase = "failed"
            return False

    def _compose_summary(self) -> str:
        """Compose the final summary."""
        exp_summary = (
            self._experiment_workflow.final_summary if self._experiment_workflow else "N/A"
        )
        current_revision = (
            self._experiment_workflow.current_revision if self._experiment_workflow else 0
        )

        paper_section = ""
        if self.run_paper_writing and self._writing_workflow is not None:
            paper_section = f"""

---

## Paper

{self.paper_writing_summary or "(no writing report available)"}

**Final LaTeX**: {self.paper_final_tex_path or "N/A"}
**Final PDF**: {self.paper_final_pdf_path or "N/A"}
"""

        return f"""# Full SciDER Workflow Summary

## Data Analysis

{self.data_summary}

---

## Experiment Results

{exp_summary}{paper_section}

---

## Workflow Metadata

- **Data Source**: {f"data_path={self.data_path}" if self.data_path else f"feature_desc={self.feature_desc!r} ({self.num_rows} rows)"}
- **Workspace**: {self.workspace_path}
- **Repo Source**: {self.repo_source or 'Not specified'}
- **Final Status**: {self.final_status}
- **Total Revisions**: {current_revision}
"""

    def _finalize(self):
        """Finalize the workflow."""
        logger.info("Finalizing workflow")

        if self.current_phase == "failed":
            self.final_summary = f"# Workflow Failed\n\nError: {self.error_message}"
        elif not self.final_summary:
            self.final_summary = "# Workflow Completed\n\nNo summary available."

        logger.info(get_separator())
        logger.info(f"Workflow completed: {self.final_status}")
        logger.info(get_separator())

    def save_summary(self, path: str | Path | None = None) -> Path:
        """Save the final summary to a file."""
        if path is None:
            path = self.workspace_path / "workflow_summary.md"
        path = Path(path)
        path.write_text(self.final_summary)
        logger.info(f"Summary saved to {path}")
        return path


def run_full_workflow(
    workspace_path: str | Path,
    user_query: str,
    data_path: str | Path | None = None,
    feature_desc: str | None = None,
    num_rows: int = 1000,
    repo_source: str | None = None,
    max_revisions: int = 3,
    data_agent_recursion_limit: int = 100,
    experiment_agent_recursion_limit: int = 100,
    data_desc: str | None = None,
    user_approval_enabled: bool = False,
    run_paper_writing: bool = False,
    paper_workspace_path: str | Path | None = None,
    paper_template_dir_path: str | Path | None = None,
    paper_template_tex_path: str | Path | None = None,
    paper_conference_guidelines_path: str | Path | None = None,
    paper_agent_recursion_limit: int = 150,
    workspace_init_config: WorkspaceInitConfig | None = None,
) -> FullWorkflow:
    """
    Convenience function to run the full SciDER workflow.

    Args:
        data_path: Path to the data file or directory to analyze
        workspace_path: Workspace directory for the experiment
        user_query: User's experiment objective
        repo_source: Optional repository source (local path or git URL)
        max_revisions: Maximum revision loops for experiment agent
        data_agent_recursion_limit: Recursion limit for DataAgent (default=100)
        experiment_agent_recursion_limit: Recursion limit for ExperimentAgent (default=100)
        data_desc: Optional additional description of the data

    Returns:
        FullWorkflow: Completed workflow with results

    Example:
        >>> result = run_full_workflow(
        ...     data_path="data/data.csv",
        ...     workspace_path="workspace",
        ...     user_query="Train an SVR model to predict prices",
        ... )
        >>> print(result.final_summary)

    Note:

        These directories are then passed to DataWorkflow and ExperimentWorkflow.
    """
    workflow = FullWorkflow(
        data_path=Path(data_path) if data_path else None,
        feature_desc=feature_desc,
        num_rows=num_rows,
        workspace_path=workspace_path,
        user_query=user_query,
        repo_source=repo_source,
        max_revisions=max_revisions,
        data_agent_recursion_limit=data_agent_recursion_limit,
        experiment_agent_recursion_limit=experiment_agent_recursion_limit,
        data_desc=data_desc,
        run_paper_writing=run_paper_writing,
        paper_workspace_path=Path(paper_workspace_path) if paper_workspace_path else None,
        paper_template_dir_path=Path(paper_template_dir_path) if paper_template_dir_path else None,
        paper_template_tex_path=Path(paper_template_tex_path) if paper_template_tex_path else None,
        paper_conference_guidelines_path=(
            Path(paper_conference_guidelines_path) if paper_conference_guidelines_path else None
        ),
        paper_agent_recursion_limit=paper_agent_recursion_limit,
        workspace_init_config=workspace_init_config,
    )
    with override_user_approval(user_approval_enabled):
        return workflow.run()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Full SciDER Workflow - Run complete workflow (DataAgent -> ExperimentAgent)",
        prog="python -m scider.workflows.full_workflow",
    )
    parser.add_argument("workspace_path", help="Workspace directory for the experiment")
    parser.add_argument("user_query", help="User's experiment objective")
    parser.add_argument(
        "--data-path",
        default=None,
        help="Path to a real data file or directory (mutually exclusive with --feature-desc)",
    )
    parser.add_argument(
        "--feature-desc",
        default=None,
        help="Description for synthetic data generation (mutually exclusive with --data-path)",
    )
    parser.add_argument(
        "--num-rows",
        type=int,
        default=1000,
        help="Number of rows to generate in hypo mode (default: 1000)",
    )
    parser.add_argument(
        "--repo-source",
        default=None,
        help="Optional repository source (local path or git URL)",
    )
    parser.add_argument(
        "--max-revisions",
        type=int,
        default=5,
        help="Maximum revision loops for ExperimentAgent (default: 5)",
    )
    parser.add_argument(
        "--data-recursion-limit",
        type=int,
        default=100,
        help="Recursion limit for DataAgent (default: 100)",
    )
    parser.add_argument(
        "--experiment-recursion-limit",
        type=int,
        default=100,
        help="Recursion limit for ExperimentAgent (default: 100)",
    )
    parser.add_argument(
        "--session-name",
        default=None,
        help="Custom session name (otherwise uses timestamp)",
    )
    parser.add_argument(
        "--data-desc",
        default=None,
        help="Optional additional description of the data",
    )

    args = parser.parse_args()

    result = run_full_workflow(
        workspace_path=args.workspace_path,
        user_query=args.user_query,
        data_path=args.data_path,
        feature_desc=args.feature_desc,
        num_rows=args.num_rows,
        repo_source=args.repo_source,
        max_revisions=args.max_revisions,
        data_agent_recursion_limit=args.data_recursion_limit,
        experiment_agent_recursion_limit=args.experiment_recursion_limit,
        data_desc=args.data_desc,
    )

    print("\n" + get_separator())
    print("FULL WORKFLOW COMPLETE")
    print(get_separator())
    print(f"\nStatus: {result.final_status}")
    print(f"\nFinal Summary:\n{result.final_summary}")