import json
import os
from datetime import datetime
from pathlib import Path

from biomni.agent import A1


PROJECT_ROOT = Path("/225040511/project/Biomni")
RESULTS_DIR = PROJECT_ROOT / "paper_repro_foxj1" / "results"
MCP_CONFIG_PATH = PROJECT_ROOT / "strata_mcp_config_shim.yaml"
QUERY = """Use the processed GEO files in
/225040511/project/Biomni/paper_repro_foxj1/data

Only do these four things:
1. data loading and schema inspection
2. reproduce the core differential expression findings from resistant vs sensitive comparisons
3. run GSEA / pathway enrichment for the LuCaP35CR resistant signature
4. partially reproduce the FOXJ1 overexpression model from the processed differential expression file

Do not perform external cohort analysis, DepMap analysis, or LP-WGS analysis.
Save all intermediate and final outputs under
/225040511/project/Biomni/paper_repro_foxj1/results

In the final answer, summarize results by these four sections and include output file paths."""


def save_run_outputs(log_entries: list[str], answer: str) -> dict[str, str]:
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_dir = RESULTS_DIR / f"main_run_{timestamp}"
    run_dir.mkdir(parents=True, exist_ok=True)

    log_txt_path = run_dir / "main_execution_log.txt"
    log_json_path = run_dir / "main_execution_log.json"
    answer_path = run_dir / "main_final_answer.txt"
    metadata_path = run_dir / "main_run_metadata.json"

    log_txt_path.write_text("\n\n".join(str(entry) for entry in log_entries), encoding="utf-8")
    log_json_path.write_text(
        json.dumps({"log_entries": log_entries}, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    answer_path.write_text(answer, encoding="utf-8")
    metadata_path.write_text(
        json.dumps(
            {
                "timestamp": timestamp,
                "mcp_config_path": str(MCP_CONFIG_PATH),
                "results_dir": str(RESULTS_DIR),
                "run_dir": str(run_dir),
                "query": QUERY,
                "log_entry_count": len(log_entries),
                "log_txt_path": str(log_txt_path),
                "log_json_path": str(log_json_path),
                "answer_path": str(answer_path),
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )

    return {
        "run_dir": str(run_dir),
        "log_txt": str(log_txt_path),
        "log_json": str(log_json_path),
        "answer_txt": str(answer_path),
        "metadata_json": str(metadata_path),
    }


agent = A1(
    path="./data",
    llm="deepseek-chat",
    source="Custom",
    base_url="https://api.deepseek.com/v1",
    api_key=os.getenv("DEEPSEEK_API_KEY", "EMPTY"),
    expected_data_lake_files=[],
    rewrite_user_query=True,
    dynamic_mcp_registration=True,
    mcp_server_top_k=20,
    mcp_tool_top_k=12,
)

# 这里传你那份包含大量 MCP servers 的总配置
agent.attach_mcp_catalog(MCP_CONFIG_PATH)

log, answer = agent.go(QUERY)
artifact_paths = save_run_outputs(log, answer)

print("\n===== SAVED ARTIFACTS =====")
for name, path in artifact_paths.items():
    print(f"{name}: {path}")

print("\n===== FINAL ANSWER =====\n")
print(answer)