from pathlib import Path
from rich.console import Console
from lobster.core.data_manager_v2 import DataManagerV2

# Create data manager with workspace
dm = DataManagerV2(
    default_backend="h5ad",                # Storage backend: "h5ad" (default) or "mudata"
    workspace_path=Path("./my_analysis"),  # Optional: workspace directory
    enable_provenance=True,                # Enable W3C-PROV tracking
    console=None,                          # Optional: Rich Console for progress
    auto_scan=True                         # Auto-load existing datasets
)

# With custom console
console = Console()
dm = DataManagerV2(
    default_backend="h5ad",
    workspace_path=Path("./my_analysis"),
    console=console
)

# Load from H5AD file
adata = dm.load_modality(
    name="my_dataset",
    source="/path/to/data.h5ad",
    adapter="transcriptomics_single_cell"
)

# Load from CSV
adata = dm.load_modality(
    name="counts_matrix",
    source="/path/to/counts.csv",
    adapter="transcriptomics_bulk"
)

# Load from GEO (via download queue)
# Typically handled by research_agent + data_expert workflow

# Get a specific modality
adata = dm.get_modality("my_dataset")

# List all loaded modalities
modalities = dm.list_modalities()
print(f"Loaded: {modalities}")

# Check shape
print(f"Cells: {adata.n_obs}, Genes: {adata.n_vars}")

# Store with lineage tracking (recommended)
dm.store_modality(
    name="my_dataset_filtered",
    adata=filtered_adata,
    parent_name="my_dataset",
    step_summary="Filtered 15% low-quality cells"
)

# Manual storage (less tracking)
dm.modalities["my_dataset_processed"] = processed_adata

# Save to H5AD (default backend)
path = dm.save_modality("my_dataset", "my_dataset.h5ad")

# Save with specific backend
path = dm.save_modality(
    "my_dataset",
    "my_dataset.h5mu",
    backend="mudata"
)

# From file path
dm.load_modality("dataset", "/path/to/file.h5ad", adapter="transcriptomics_single_cell")

# From pandas DataFrame
import pandas as pd
df = pd.read_csv("counts.csv", index_col=0)
dm.load_modality("from_csv", df, adapter="transcriptomics_bulk")

# From existing AnnData
import anndata
adata = anndata.read_h5ad("existing.h5ad")
dm.load_modality("imported", adata, adapter="transcriptomics_single_cell")

geo_gse12345                      # Raw loaded data
├─ geo_gse12345_quality_assessed  # After QC metrics
├─ geo_gse12345_filtered          # After cell/gene filtering
├─ geo_gse12345_normalized        # After normalization
├─ geo_gse12345_clustered         # After clustering
└─ geo_gse12345_annotated         # After cell type annotation

# Get modalities with full lineage information
modalities = dm.list_modalities_with_lineage()

for mod in modalities:
    print(f"{mod['name']}: v{mod['version']}")
    print(f"  Parent: {mod.get('parent', 'None')}")
    print(f"  Step: {mod.get('processing_step', 'raw')}")

from lobster.core.interfaces.backend import IDataBackend

class MyS3Backend(IDataBackend):
    def save(self, adata, path, **kwargs): ...
    def load(self, path, **kwargs): ...

dm.register_backend("s3", MyS3Backend(bucket="my-bucket"))

# Save with MuData backend
dm.save_modality("integrated", "integrated.h5mu", backend="mudata")

from lobster.core.interfaces.adapter import IModalityAdapter

class MyAdapter(IModalityAdapter):
    def from_source(self, source, **kwargs): ...
    def validate(self, adata, strict=False): ...
    def standardize(self, adata): ...

dm.register_adapter("my_format", MyAdapter())

def create_transcriptomics_expert(
    data_manager: DataManagerV2,
    callback_handler,
    agent_name: str = "transcriptomics_expert",
    delegation_tools: list = None,
    workspace_path = None
):
    """Factory creates agent with data manager access."""
    # Agent tools can access dm.modalities, dm.get_modality(), etc.

@tool
def filter_cells(modality_name: str, min_genes: int = 200) -> str:
    """Filter low-quality cells from dataset."""
    adata = data_manager.get_modality(modality_name)

    # Call service
    result, stats, ir = filtering_service.filter(adata, min_genes=min_genes)

    # Store with lineage
    output_name = f"{modality_name}_filtered"
    data_manager.store_modality(
        name=output_name,
        adata=result,
        parent_name=modality_name,
        step_summary=f"Filtered cells with <{min_genes} genes"
    )

    # Log to provenance
    data_manager.log_tool_usage(
        tool_name="filter_cells",
        parameters={"modality": modality_name, "min_genes": min_genes},
        ir=ir
    )

    return f"Filtered: {stats['cells_removed']} cells removed"

.lobster_workspace/
├── data/                    # Saved H5AD/H5MU files
│   ├── geo_gse12345.h5ad
│   └── geo_gse12345_filtered.h5ad
├── exports/                 # Exported results
│   └── markers.csv
├── cache/                   # Temporary cache
├── literature_cache/        # Publication data
│   ├── publications/
│   └── parsed_docs/
├── .lobster/
│   └── queues/             # Download/publication queues
│       ├── download_queue.jsonl
│       └── publication_queue.jsonl
├── plots/                   # Saved visualizations
├── notebooks/               # Exported pipelines
└── .session.json            # Session metadata

# Session 1
dm = DataManagerV2(workspace_path="./my_analysis")
dm.load_modality("dataset", "data.h5ad", adapter="transcriptomics_single_cell")
dm.save_modality("dataset", "dataset.h5ad")
# Exit session

# Session 2 - dataset auto-loaded
dm = DataManagerV2(workspace_path="./my_analysis", auto_scan=True)
adata = dm.get_modality("dataset")  # Available immediately

# Log tool usage with IR for reproducibility
dm.log_tool_usage(
    tool_name="normalize",
    parameters={"method": "log1p", "target_sum": 10000},
    description="Normalize counts per cell",
    ir=analysis_step  # AnalysisStep for notebook export
)

# Access provenance data
activities = dm.provenance.get_all_activities()

# Export current session
path = dm.export_notebook(
    name="qc_workflow",
    description="Quality control for 10X data",
    filter_strategy="successful"  # Only successful operations
)

# Run on new data
result = dm.run_notebook(
    notebook_path="qc_workflow.ipynb",
    input_modality="new_dataset",
    parameters={"min_genes": 300}
)