streamlit-client/app.py

import json
import os
import shutil
import sys
import tempfile
import time
import zipfile
from collections import defaultdict
from datetime import datetime
from pathlib import Path

import streamlit as st

os.environ["CODING_AGENT_VERSION"] = "v3"
os.environ.setdefault("SCIEVO_ENABLE_OPENHANDS", "0")
# Disable Reasoning Bank when running workflows from Streamlit client
os.environ["REASONING_BANK_ENABLED"] = "false"

sys.path.insert(0, str(Path(__file__).parent.parent))

from scievo.agents import ideation_agent
from scievo.agents.ideation_agent.state import IdeationAgentState
from scievo.core.brain import Brain
from scievo.core.llms import ModelRegistry
from scievo.workflows.data_workflow import DataWorkflow
from scievo.workflows.experiment_workflow import ExperimentWorkflow
from scievo.workflows.full_workflow_with_ideation import FullWorkflowWithIdeation

try:
    from streamlit_file_browser import st_file_browser
except ImportError:
    st_file_browser = None

st.set_page_config(page_title="SciDER Chat", layout="centered")

def get_theme_css() -> str:
    """Return CSS for light theme - ensure all UI components use light mode."""
    return """
    <style>
    /* Force light color scheme - overrides browser/system dark mode */
    html, :root { color-scheme: light !important; }
    /* Base backgrounds */
    [data-testid="stApp"], .stApp { background-color: #ffffff !important; }
    [data-testid="stAppViewContainer"] { background-color: #ffffff !important; }
    [data-testid="stVerticalBlock"], [data-testid="block-container"] { background-color: #ffffff !important; }
    [data-testid="stChatMessage"] { background-color: transparent !important; border: none !important; box-shadow: none !important; }
    [data-testid="stSidebar"] { background-color: #ffffff !important; }
    [data-testid="stExpander"], .stForm { background-color: #ffffff !important; }
    /* Typography */
    h1, h2, h3, h4, h5, h6 { color: #384166 !important; }
    [data-testid="stChatMessage"] h1, [data-testid="stChatMessage"] h2,
    [data-testid="stChatMessage"] h3, [data-testid="stChatMessage"] h4,
    [data-testid="stChatMessage"] h5, [data-testid="stChatMessage"] h6 { color: inherit !important; }
    .stMarkdown, p, span, label, [data-testid="stMarkdown"], [data-testid="stCaptionContainer"],
    [data-testid="stVerticalBlock"] { color: #262730 !important; }
    [data-testid="stChatMessage"] .stMarkdown, [data-testid="stChatMessage"] p,
    [data-testid="stChatMessage"] span { color: #262730 !important; }
    /* Inputs - light mode */
    input, textarea { color: #262730 !important; background-color: #ffffff !important; border-color: #d1d5db !important; }
    input::placeholder, textarea::placeholder { color: #6b7280 !important; }
    div[data-baseweb="input"] input, div[data-baseweb="input"] { background-color: #ffffff !important; color: #262730 !important; border-color: #d1d5db !important; }
    /* Password input eye icon - light mode, soft gray */
    [title="Show password text"], [title="Hide password text"],
    div[data-baseweb="input"] button, div[data-baseweb="input"] [role="button"],
    div[data-baseweb="input"] [data-baseweb="button"] {
        color: #6b7280 !important; background-color: transparent !important;
    }
    div[data-baseweb="input"] svg, div[data-baseweb="input"] path,
    div[data-baseweb="input"] svg path {
        fill: #6b7280 !important; color: #6b7280 !important;
    }
    div[data-baseweb="select"] > div, div[role="combobox"] { background-color: #ffffff !important; color: #262730 !important; border-color: #d1d5db !important; }
    div[role="listbox"], div[role="listbox"] li { background-color: #ffffff !important; color: #262730 !important; }
    /* All Streamlit/BaseWeb buttons - force light mode */
    [data-testid="stButton"] button,
    [data-testid="stButton"] > div,
    [data-testid="stButton"] > div > div,
    div[data-baseweb="button"],
    div[data-baseweb="button"] button,
    div[data-baseweb="button"] > div,
    div[data-baseweb="button"] span,
    button[kind="secondary"],
    button[kind="tertiary"] {
        background-color: #f0f2f6 !important;
        color: #262730 !important;
        border-color: #d1d5db !important;
    }
    [data-testid="stButton"]:hover button,
    [data-testid="stButton"]:hover > div,
    div[data-baseweb="button"]:hover,
    div[data-baseweb="button"]:hover > div {
        background-color: #e8eaed !important;
        color: #262730 !important;
    }
    /* Primary button (Save API Keys) - keep theme color */
    button[data-baseweb="primary"],
    div[data-baseweb="button"][data-pseudo="-primary"],
    [data-testid="stButton"]:has(button[data-baseweb="primary"]) button {
        background-color: #384166 !important;
        color: #ffffff !important;
        border-color: #384166 !important;
    }
    /* Alerts - light mode (info, warning, error) */
    .stAlert, [data-testid="stAlert"], [data-baseweb="notification"] { 
        background-color: #eff6ff !important; color: #1e40af !important; 
        border: 1px solid #93c5fd !important; 
    }
    /* File uploader - light mode */
    [data-testid="stFileUploader"],
    [data-testid="stFileUploader"] section,
    [data-testid="stFileUploader"] div,
    [data-testid="stFileUploader"] span,
    [data-testid="stFileUploader"] label,
    [data-testid="stFileUploader"] * {
        background-color: #ffffff !important;
        color: #262730 !important;
        border-color: #d1d5db !important;
    }
    [data-testid="stFileUploader"] [data-baseweb="file-uploader"],
    [data-testid="stFileUploader"] [data-baseweb="fileuploader"] {
        background-color: #f8f9fa !important;
        border: 2px dashed #d1d5db !important;
    }
    /* Expanders */
    [data-testid="stExpander"] details, [data-testid="stExpander"] summary { color: #262730 !important; background-color: #ffffff !important; }
    /* Code blocks (LLM output, markdown) - light mode */
    pre, code, [data-testid="stMarkdown"] pre, [data-testid="stMarkdown"] code,
    [data-testid="stChatMessage"] pre, [data-testid="stChatMessage"] code,
    .stMarkdown pre, .stMarkdown code, .highlight, .hljs,
    pre code, .highlight pre, [data-testid="stCodeBlock"] {
        background-color: #f8f9fa !important;
        color: #262730 !important;
        border: 1px solid #e5e7eb !important;
    }
    [data-testid="stChatMessage"] pre, [data-testid="stChatMessage"] code,
    [data-testid="stChatMessage"] .highlight, [data-testid="stChatMessage"] .hljs,
    [data-testid="stChatMessage"] pre code { color: #262730 !important; background-color: #f8f9fa !important; }
    /* Syntax highlight - light theme token colors */
    .hljs-keyword, .hljs-selector-tag { color: #7c3aed !important; }
    .hljs-string, .hljs-attr { color: #0d9488 !important; }
    .hljs-number { color: #dc2626 !important; }
    .hljs-comment { color: #6b7280 !important; }
    .hljs-title, .hljs-function { color: #2563eb !important; }
    </style>
    """


st.markdown(get_theme_css(), unsafe_allow_html=True)


def register_all_models(user_api_key=None, user_model=None):
    api_key = user_api_key or os.getenv("GEMINI_API_KEY") or os.getenv("OPENAI_API_KEY")
    if not api_key:
        return False

    default_model = user_model or os.getenv("SCIEVO_DEFAULT_MODEL", "gemini/gemini-2.5-flash-lite")
    openai_api_key = (
        user_api_key
        if user_api_key and "openai" in default_model.lower()
        else os.getenv("OPENAI_API_KEY")
    )

    models = [
        ("ideation", default_model, api_key),
        ("data", default_model, api_key),
        ("plan", default_model, api_key),
        ("history", default_model, api_key),
        ("experiment_agent", default_model, api_key),
        ("experiment_coding", default_model, api_key),
        ("experiment_execute", default_model, api_key),
        ("experiment_summary", default_model, api_key),
        ("experiment_monitor", default_model, api_key),
        ("paper_search", default_model, api_key),
        ("metric_search", default_model, api_key),
        ("critic", default_model, api_key),
        ("mem", default_model, api_key),
    ]

    embed_model = os.getenv("EMBED_MODEL", "text-embedding-004")
    embed_api_key = os.getenv("EMBED_API_KEY", openai_api_key or api_key)
    models.append(("embed", embed_model, embed_api_key))

    for name, model, key in models:
        ModelRegistry.register(name=name, model=model, api_key=key)

    return True


def stream_markdown(text, delay=0.02):
    buf = ""
    slot = st.empty()
    for line in text.split("\n"):
        buf += line + "\n"
        slot.markdown(buf)
        time.sleep(delay)


def render_intermediate_state(intermediate_state):
    if not intermediate_state:
        return
    by_node = defaultdict(list)
    for item in intermediate_state:
        by_node[item.get("node_name", "unknown")].append(item.get("output", ""))

    st.divider()
    st.subheader("Intermediate States")
    for node, outputs in by_node.items():
        with st.expander(node, expanded=False):
            for i, content in enumerate(outputs, 1):
                st.markdown(f"**Step {i}**")
                st.markdown(content)


def run_ideation(q):
    s = IdeationAgentState(user_query=q)
    r = st.session_state.ideation_graph.invoke(s, {"recursion_limit": 50})
    rs = IdeationAgentState(**r)
    out = []
    if rs.output_summary:
        out.append("## Research Ideas Summary\n\n" + rs.output_summary)
    if rs.novelty_score is not None:
        out.append(
            "## Novelty Evaluation\n```json\n"
            + json.dumps(
                {
                    "novelty_score": rs.novelty_score,
                    "feedback": rs.novelty_feedback,
                },
                indent=2,
            )
            + "\n```"
        )
    if rs.research_ideas:
        out.append("## Generated Research Ideas\n")
        for i, idea in enumerate(rs.research_ideas[:5], 1):
            out.append(f"### {i}. {idea.get('title','')}\n{idea.get('description','')}")
    return ("\n\n".join(out) if out else "No result", rs.intermediate_state)


def run_data(path, q):
    # Ensure path is absolute and exists
    data_path = Path(path).resolve()
    if not data_path.exists():
        return f"Error: Data path does not exist: {data_path}", []

    # Log path for debugging
    logger = __import__("loguru").logger
    logger.info(f"Running data analysis on path: {data_path}")
    logger.info(
        f"Path exists: {data_path.exists()}, is_dir: {data_path.is_dir()}, is_file: {data_path.is_file()}"
    )

    w = DataWorkflow(
        data_path=data_path,
        workspace_path=st.session_state.workspace_path,
        recursion_limit=100,
    )
    w.run()
    intermediate_state = getattr(w, "data_agent_intermediate_state", [])
    if w.final_status != "success":
        error_msg = w.error_message or "Data workflow failed"
        return f"Data workflow failed: {error_msg}", intermediate_state
    out = ["## Data Analysis Complete"]
    if w.data_summary:
        out.append(w.data_summary)
    return "\n\n".join(out), intermediate_state


def run_experiment(q, path):
    if path:
        # Ensure path is absolute and exists
        analysis_path = Path(path).resolve()
        if not analysis_path.exists():
            return f"Error: Data analysis file does not exist: {analysis_path}", []

        logger = __import__("loguru").logger
        logger.info(f"Running experiment with analysis file: {analysis_path}")
        logger.info(f"Path exists: {analysis_path.exists()}, is_file: {analysis_path.is_file()}")

        w = ExperimentWorkflow.from_data_analysis_file(
            workspace_path=st.session_state.workspace_path,
            user_query=q,
            data_analysis_path=str(analysis_path),
            max_revisions=5,
            recursion_limit=100,
        )
    else:
        return "No data analysis file", []
    w.run()
    return w.final_summary or "Experiment finished", w.experiment_agent_intermediate_state


def run_full(cfg):
    data_path = None
    if cfg.get("data_path"):
        data_path = Path(cfg["data_path"]).resolve()
        if not data_path.exists():
            return f"Error: Data path does not exist: {data_path}", []

    logger = __import__("loguru").logger
    if data_path:
        logger.info(f"Running full workflow with data path: {data_path}")
        logger.info(
            f"Path exists: {data_path.exists()}, is_dir: {data_path.is_dir()}, is_file: {data_path.is_file()}"
        )

    w = FullWorkflowWithIdeation(
        user_query=cfg["query"],
        workspace_path=st.session_state.workspace_path,
        data_path=data_path,
        run_data_workflow=cfg["run_data"],
        run_experiment_workflow=cfg["run_exp"],
        max_revisions=5,
    )
    w.run()

    # Aggregate intermediate state from all phases for subagent output display
    aggregated = []
    for item in w.ideation_intermediate_state:
        aggregated.append({
            "node_name": f"ideation/{item.get('node_name', 'unknown')}",
            "output": item.get("output", ""),
        })
    if w._data_workflow:
        for item in getattr(w._data_workflow, "data_agent_intermediate_state", []) or []:
            aggregated.append({
                "node_name": f"data/{item.get('node_name', 'unknown')}",
                "output": item.get("output", ""),
            })
    if w._experiment_workflow:
        for item in getattr(w._experiment_workflow, "experiment_agent_intermediate_state", []) or []:
            aggregated.append({
                "node_name": f"experiment/{item.get('node_name', 'unknown')}",
                "output": item.get("output", ""),
            })
    return w.final_summary or "Workflow finished", aggregated


def get_upload_temp_dir() -> Path:
    """Return temp directory for uploaded files. Clean old dirs on startup."""
    base = Path(tempfile.gettempdir()) / "scider_uploads"
    base.mkdir(parents=True, exist_ok=True)
    # Clean dirs older than 1 hour (handles closed sessions)
    now = time.time()
    for d in base.iterdir():
        if d.is_dir() and (now - d.stat().st_mtime) > 3600:
            try:
                shutil.rmtree(d)
            except OSError:
                pass
    return base


def save_and_extract_upload(uploaded_file) -> Path | None:
    """Save uploaded zip to temp dir, extract it, return path to extracted dir."""
    if uploaded_file is None or not uploaded_file.name.lower().endswith(".zip"):
        return None
    base = get_upload_temp_dir()
    dest_dir = Path(tempfile.mkdtemp(dir=base))
    zip_path = dest_dir / uploaded_file.name
    with open(zip_path, "wb") as f:
        f.write(uploaded_file.getvalue())
    extract_dir = dest_dir / "extracted"
    extract_dir.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(zip_path, "r") as zf:
        zf.extractall(extract_dir)
    zip_path.unlink()
    # Return absolute path to ensure it works in container environments
    return extract_dir.resolve()


def find_data_analysis_file(extract_dir: Path) -> Path | None:
    """Find data_analysis.md in extracted dir (root or first subdir)."""
    candidates = [extract_dir / "data_analysis.md", extract_dir / "analysis.md"]
    for c in candidates:
        if c.exists():
            return c
    for p in extract_dir.rglob("data_analysis.md"):
        return p
    for p in extract_dir.rglob("analysis.md"):
        return p
    return None


def _rm_upload_root(p: Path):
    """Remove the scider_uploads session dir (go up to find it)."""
    cur = Path(p).resolve().parent if Path(p).resolve().is_file() else Path(p).resolve()
    while cur != cur.parent:
        parent = cur.parent
        if parent.name == "scider_uploads":
            try:
                shutil.rmtree(cur)
            except OSError:
                pass
            return
        cur = parent


def cleanup_uploaded_data():
    """Remove temp uploaded data and restore workspace_path to default."""
    for key in ("uploaded_data_path", "uploaded_experiment_path", "uploaded_full_data_path"):
        path = st.session_state.get(key)
        if path and isinstance(path, (str, Path)):
            _rm_upload_root(Path(path))
            if key in st.session_state:
                del st.session_state[key]
    # Restore agent workspace to default
    if "default_workspace_path" in st.session_state:
        st.session_state.workspace_path = st.session_state.default_workspace_path


def get_case_study_memory_paths() -> list[Path]:
    """Return list of case-study-memory directories to search (cwd and project root)."""
    paths = []
    cwd = Path.cwd()
    app_dir = Path(__file__).parent
    project_root = app_dir.parent
    for p in [cwd / "case-study-memory", project_root / "case-study-memory", app_dir / "case-study-memory"]:
        if p.exists() and p.is_dir() and p not in paths:
            paths.append(p)
    return paths


def list_available_chat_files() -> list[tuple[Path, dict]]:
    """List all chat_history.json files in case-study-memory, return (path, metadata_for_display)."""
    results = []
    seen = set()
    for base in get_case_study_memory_paths():
        dirs = [d for d in base.iterdir() if d.is_dir() and (d / "chat_history.json").exists()]
        dirs.sort(key=lambda x: x.stat().st_mtime if x.exists() else 0, reverse=True)
        for memo_dir in dirs:
            chat_file = memo_dir / "chat_history.json"
            if chat_file.exists():
                try:
                    with open(chat_file, encoding="utf-8") as f:
                        data = json.load(f)
                    ts = data.get("timestamp", "")[:19].replace("T", " ")
                    wf = data.get("workflow_type", "unknown")
                    meta = data.get("metadata", {})
                    q = meta.get("query") or ""
                    # Only show query (no data_path, timestamp, workflow_type, etc.)
                    query = (q[:80] + "...") if len(q) > 83 else (q or memo_dir.name)
                    key = str(chat_file.resolve())
                    if key not in seen:
                        seen.add(key)
                        results.append((chat_file, {"label": query, "timestamp": ts, "workflow_type": wf}))
                except Exception:
                    results.append((chat_file, {"label": f"{memo_dir.name} | (parse error)", "timestamp": "", "workflow_type": "unknown"}))
    return results


def load_chat_from_file(chat_path: Path) -> list[dict]:
    """Load messages from chat_history.json. Returns list of message dicts."""
    with open(chat_path, encoding="utf-8") as f:
        data = json.load(f)
    return data.get("messages", [])


def _render_case_file_preview(event: dict | None, root_path: Path):
    """Render a light-theme preview for selected files from streamlit-file-browser."""
    if not event or not isinstance(event, dict):
        return
    if event.get("type") != "SELECT_FILE":
        return
    target = (event.get("target") or {}).get("path")
    if not target:
        return
    file_path = (root_path / target).resolve()
    if not file_path.exists() or file_path.is_dir():
        st.warning(f"File `{target}` not found under current workspace.")
        return

    st.markdown("#### File Preview")
    suffix = file_path.suffix.lower()
    if suffix in {".png", ".jpg", ".jpeg", ".gif", ".webp"}:
        st.image(str(file_path))
        return

    # Text/code preview (light theme via existing CSS on st.code)
    max_chars = 300_000
    with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
        content = f.read(max_chars + 1)
    if len(content) > max_chars:
        content = content[:max_chars] + "\n\n... [truncated]"

    lang_map = {
        ".py": "python",
        ".js": "javascript",
        ".ts": "typescript",
        ".json": "json",
        ".md": "markdown",
        ".sh": "bash",
        ".yml": "yaml",
        ".yaml": "yaml",
        ".csv": "text",
        ".txt": "text",
    }
    st.code(content, language=lang_map.get(suffix, "text"))


def get_next_memo_number(memory_dir: Path) -> int:
    if not memory_dir.exists():
        return 1

    existing_memos = [
        d.name for d in memory_dir.iterdir() if d.is_dir() and d.name.startswith("memo_")
    ]

    if not existing_memos:
        return 1

    numbers = []
    for memo in existing_memos:
        try:
            num = int(memo.replace("memo_", ""))
            numbers.append(num)
        except ValueError:
            continue

    return max(numbers) + 1 if numbers else 1


def _case_study_base() -> Path:
    """Return case-study-memory base dir. Supports both local and Docker layouts."""
    # Local: case-study-memory at project root (sibling of streamlit-client)
    local = Path(__file__).parent.parent / "case-study-memory"
    # Docker: case-study-memory mounted at streamlit-client/case-study-memory
    docker = Path(__file__).parent / "case-study-memory"
    # Prefer the one that has memo workspaces (subdirs with workspace/)
    for candidate in (local, docker):
        if candidate.exists():
            has_workspace = sum(1 for d in candidate.iterdir() if d.is_dir() and (d / "workspace").exists())
            if has_workspace > 0:
                return candidate
    return local if local.exists() else docker


def allocate_memo_workspace() -> tuple[Path, Path]:
    """Allocate memo_X and memo_X/workspace for this run. Returns (memo_dir, workspace_path)."""
    base_dir = _case_study_base()
    base_dir.mkdir(parents=True, exist_ok=True)
    memo_number = get_next_memo_number(base_dir)
    memo_dir = base_dir / f"memo_{memo_number}"
    memo_dir.mkdir(parents=True, exist_ok=True)
    workspace_path = memo_dir / "workspace"
    workspace_path.mkdir(parents=True, exist_ok=True)
    return memo_dir, workspace_path


def save_chat_history(
    messages: list, workflow_type: str, metadata: dict = None, memo_dir: Path | None = None
):
    """Save chat history. If memo_dir is provided, save there; else allocate new memo."""
    base_dir = _case_study_base()
    base_dir.mkdir(parents=True, exist_ok=True)

    if memo_dir is None:
        memo_number = get_next_memo_number(base_dir)
        memo_dir = base_dir / f"memo_{memo_number}"
        memo_dir.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now().isoformat()

    chat_data = {
        "timestamp": timestamp,
        "workflow_type": workflow_type,
        "metadata": metadata or {},
        "messages": messages,
    }

    chat_file = memo_dir / "chat_history.json"
    with open(chat_file, "w", encoding="utf-8") as f:
        json.dump(chat_data, f, indent=2, ensure_ascii=False)

    return memo_dir


if "api_key" not in st.session_state:
    st.session_state.api_key = os.getenv("GEMINI_API_KEY") or os.getenv("OPENAI_API_KEY") or ""
if "anthropic_api_key" not in st.session_state:
    st.session_state.anthropic_api_key = os.getenv("ANTHROPIC_API_KEY") or ""
if "default_model" not in st.session_state:
    st.session_state.default_model = os.getenv(
        "SCIEVO_DEFAULT_MODEL", "gemini/gemini-2.5-flash-lite"
    )
if "view_mode" not in st.session_state:
    st.session_state.view_mode = "live"  # "live" | "case_study"

# Case Study mode: no API key required, load from case-study-memory, same UI as live
if st.session_state.view_mode == "case_study":
    # Same header as live mode
    col_title, col_reset = st.columns([5, 1])
    with col_title:
        st.title("SciDER Research Assistant")
    with col_reset:
        if st.button("← Live", help="Back to live chat", key="back_to_live"):
            st.session_state.view_mode = "live"
            st.rerun()

    available = list_available_chat_files()
    if not available:
        st.info("No case studies found in case-study-memory. Run a workflow first to save conversations.")
        st.stop()

    options = [m["label"] for _, m in available]
    paths = [p for p, _ in available]
    default_idx = next((i for i, p in enumerate(paths) if "Kepler_Exoplanets_Prediction" in str(p)), 0)
    idx = st.selectbox("Select case study", range(len(options)), format_func=lambda i: options[i], index=default_idx, key="casestudy_select")
    selected_path = paths[idx]
    loaded = load_chat_from_file(selected_path)
    st.divider()

    for m in loaded:
        with st.chat_message(m["role"]):
            st.markdown(m["content"])
            if m.get("intermediate_state"):
                render_intermediate_state(m["intermediate_state"])

    # File browser for case study mode — browse workspace of selected case
    case_dir = selected_path.parent
    ws_dir = case_dir / "workspace"
    browse_path = ws_dir if ws_dir.exists() else case_dir
    browse_root = str(browse_path.resolve())
    current_case_key = str(case_dir.resolve())
    # Bump browser key on case switch, but avoid forced rerun (can reset component state repeatedly).
    if st.session_state.get("case_fb_selected_key") != current_case_key:
        st.session_state.case_fb_selected_key = current_case_key
        st.session_state.case_fb_nonce = st.session_state.get("case_fb_nonce", 0) + 1
    # Use a case-specific key so switching case study resets browser state correctly
    fb_key = f"case_study_file_browser_{case_dir.name}_{st.session_state.get('case_fb_nonce', 0)}"
    with st.expander("📂 Workspace Files (Agent Code)", expanded=False):
        st.caption(f"Browsing: `{browse_root}`")
        if st_file_browser is not None:
            fb_event = st_file_browser(
                browse_root,
                key=fb_key,
                show_choose_file=True,
                show_download_file=True,
                show_delete_file=False,
                show_new_folder=False,
                show_upload_file=False,
                show_preview=False,
            )
            _render_case_file_preview(fb_event, browse_path.resolve())
        else:
            st.info("Install `streamlit-file-browser` to browse workspace files in Case Study mode.")
    st.stop()

# Live mode: require API key - enhanced login page
if not st.session_state.api_key:
    st.markdown(
        """
        <div style="
            text-align: center;
            padding: 2rem 0 1.5rem;
            border-bottom: 1px solid #e5e7eb;
        ">
            <h1 style="font-size: 2rem; font-weight: 600; color: #384166; margin-bottom: 0.25rem;">🍎 SciDER</h1>
            <p style="color: #6b7280; font-size: 0.95rem;">SciDER: Scientific Data-centric End-to-end Researcher</p>
        </div>
        """,
        unsafe_allow_html=True,
    )

    # Case Study - no API key required
    st.markdown("#### Browse without API key")
    st.markdown("Explore saved conversations from previous runs — no setup required.")
    if st.button("📂 Open Case Study", help="Browse saved chat history", key="casestudy_from_login", use_container_width=True):
        st.session_state.view_mode = "case_study"
        st.rerun()

    st.markdown("---")
    st.markdown("#### Sign in to use Live Assistant")
    st.info("Enter your API keys below to run ideation, data analysis, and experiments.")

    # Model provider selection
    col1, col2 = st.columns(2)
    with col1:
        model_option = st.selectbox(
            "Model Provider",
            ["Gemini", "OpenAI"],
            index=0 if "gemini" in st.session_state.default_model.lower() else 1,
        )

    with col2:
        api_key_input = st.text_input(
            f"{model_option} API Key",
            type="password",
            placeholder=f"Enter your {model_option} API key",
            value="",
            help=f"Required for {model_option} models",
        )

    anthropic_api_key_input = st.text_input(
        "Anthropic (Claude) API Key",
        type="password",
        placeholder="Optional — for Claude coding agent",
        value="",
        help="Recommended for code generation",
    )

    st.markdown("")
    if st.button("Save API Keys", type="primary", use_container_width=True):
        if api_key_input:
            st.session_state.api_key = api_key_input
            if model_option == "Gemini":
                st.session_state.default_model = "gemini/gemini-2.5-flash-lite"
            else:
                st.session_state.default_model = "gpt-4o-mini"

            # Save Anthropic API key if provided
            if anthropic_api_key_input:
                st.session_state.anthropic_api_key = anthropic_api_key_input
                os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key_input

            st.rerun()
        else:
            st.error("Please enter a valid API key for the selected model provider")
    st.stop()

col_title, col_reset = st.columns([5, 1])
with col_title:
    st.title("SciDER Research Assistant")
with col_reset:
    if st.button("🔄 Reset", help="Clear all chat history", type="secondary"):
        cleanup_uploaded_data()
        st.session_state.messages = [
            {
                "role": "assistant",
                "content": "Hello. I can run ideation, data analysis, experiments, or a full workflow.\n\nPlease select a workflow type below to get started.",
            }
        ]
        if "selected_workflow" in st.session_state:
            st.session_state.selected_workflow = None
        # Note: API keys are preserved on reset (user doesn't need to re-enter them)
        st.rerun()

if "initialized" not in st.session_state:
    # Load environment variables from .env file
    try:
        from dotenv import load_dotenv

        # Try to load from parent directory (project root)
        env_path = Path(__file__).parent.parent / ".env"
        if env_path.exists():
            load_dotenv(env_path)
        else:
            # Fallback: try current directory
            load_dotenv()
    except Exception as e:
        logger = __import__("loguru", fromlist=["logger"]).logger
        logger.warning(f"Failed to load .env file: {e}")

    # Ensure ANTHROPIC_API_KEY is available for Claude Agent SDK
    if not os.getenv("ANTHROPIC_API_KEY"):
        # First try to get from session state (user input)
        anthropic_key = st.session_state.get("anthropic_api_key", "")
        if anthropic_key:
            os.environ["ANTHROPIC_API_KEY"] = anthropic_key
        else:
            # Fallback: try to get from user's main API key if it's an Anthropic key
            user_key = st.session_state.get("api_key", "")
            if user_key and ("anthropic" in user_key.lower() or user_key.startswith("sk-ant-")):
                os.environ["ANTHROPIC_API_KEY"] = user_key
                st.session_state.anthropic_api_key = user_key

    if not os.getenv("BRAIN_DIR"):
        os.environ["BRAIN_DIR"] = str(Path.cwd() / "tmp_brain")
    Brain()
    if register_all_models(st.session_state.api_key, st.session_state.default_model):
        st.session_state.ideation_graph = ideation_agent.build().compile()
        st.session_state.initialized = True
    else:
        st.error("Failed to register models. Please check your API key.")
        st.stop()

if "messages" not in st.session_state:
    st.session_state.messages = [
        {
            "role": "assistant",
            "content": "Hello. I can run ideation, data analysis, experiments, or a full workflow.\n\nPlease select a workflow type below to get started.",
        }
    ]

if "workspace_path" not in st.session_state:
    st.session_state.workspace_path = Path(__file__).parent.parent / "workspace"
if "default_workspace_path" not in st.session_state:
    st.session_state.default_workspace_path = Path(__file__).parent.parent / "workspace"
# If workspace_path points to expired temp upload dir, restore to default
_ws = st.session_state.workspace_path
if isinstance(_ws, (str, Path)) and "scider_uploads" in str(_ws) and not Path(_ws).exists():
    cleanup_uploaded_data()

if "selected_workflow" not in st.session_state:
    st.session_state.selected_workflow = None

# Workflow selection UI - buttons (placed at top for visibility)
st.subheader("Select Workflow Type")
col1, col2, col3, col4 = st.columns(4)

with col1:
    if st.button("💡 Ideation", use_container_width=True, key="btn_ideation"):
        st.session_state.selected_workflow = "ideation"
        st.rerun()

with col2:
    if st.button("📊 Data Analysis", use_container_width=True, key="btn_data"):
        st.session_state.selected_workflow = "data"
        st.rerun()

with col3:
    if st.button("🧪 Experiment", use_container_width=True, key="btn_experiment"):
        st.session_state.selected_workflow = "experiment"
        st.rerun()

with col4:
    if st.button("🍎 Full Workflow", use_container_width=True, key="btn_full"):
        st.session_state.selected_workflow = "full"
        st.rerun()

st.divider()

# Workspace Code File Browser — hidden from live/agent chat (shown in case_study mode only)
if st.session_state.view_mode != "live" and st_file_browser is not None:
    base = _case_study_base()
    memo_workspaces = []
    if base.exists():
        for d in sorted(base.iterdir(), key=lambda x: x.stat().st_mtime if x.exists() else 0, reverse=True):
            if not d.is_dir():
                continue
            ws = d / "workspace"
            # Include all subdirs: use workspace/ if exists, else the dir itself
            browse_path = ws if ws.exists() else d
            memo_workspaces.append((d.name, str(browse_path.resolve())))

    with st.expander("📂 Workspace Files (Agent Code)", expanded=False):
        # Use custom folder path method (same flow that worked for custom)
        if "fb_selected_memo" not in st.session_state and memo_workspaces:
            st.session_state.fb_selected_memo = memo_workspaces[0][1]

        # Quick-select memo workspaces
        memo_path = None
        if memo_workspaces:
            opts = [p for _, p in memo_workspaces]
            labels = {p: n for n, p in memo_workspaces}
            sel = st.selectbox(
                "Select memo workspace",
                options=opts,
                format_func=lambda p: labels.get(p, p),
                key="fb_memo_select",
            )
            if sel:
                st.session_state.fb_selected_memo = sel
                memo_path = sel

        # Text input for custom path (overrides memo selection when filled)
        custom = st.text_input(
            "Or enter custom folder path (browses its `workspace` subdir if present)",
            placeholder="e.g. ./case-study-memory/memo_1",
            key="fb_custom_path",
        )

        # Resolve path - same logic as custom
        if custom.strip():
            p = Path(custom.strip()).expanduser().resolve()
        elif memo_path:
            p = Path(memo_path).resolve()
        else:
            p = base.resolve() if base.exists() else Path.cwd()

        ws_sub = p / "workspace"
        if ws_sub.exists():
            ws_path = ws_sub
        elif p.exists():
            ws_path = p
        else:
            st.warning(f"Path does not exist: `{p}`")
            ws_path = base.resolve() if base.exists() else Path.cwd()

        ws_path_str = str(ws_path.resolve())
        st.caption(f"Browsing: `{ws_path_str}`")
        _fb_event = st_file_browser(
            ws_path_str,
            key="workspace_file_browser",
            show_choose_file=True,
            show_download_file=True,
            show_delete_file=False,
            show_new_folder=False,
            show_upload_file=False,
            show_preview=True,
        )
elif st.session_state.view_mode != "live":
    with st.expander("📂 Workspace Files (Agent Code)", expanded=False):
        st.info("Install `streamlit-file-browser` to browse workspace files.")
# When view_mode == "live", file browser is hidden from agent chat

for m in st.session_state.messages:
    with st.chat_message(m["role"]):
        st.markdown(m["content"])
        # Render subagent intermediate states if persisted (survives st.rerun)
        if m.get("intermediate_state"):
            render_intermediate_state(m["intermediate_state"])

# Workflow input forms
workflow_config = None

if st.session_state.selected_workflow == "ideation":
    with st.form("ideation_form", clear_on_submit=True):
        st.markdown("### 💡 Ideation Workflow")
        topic = st.text_input("Research Topic", placeholder="Enter your research topic here...")
        submitted = st.form_submit_button("Run Ideation", type="primary")
        if submitted and topic:
            workflow_config = {"type": "ideation", "query": topic}
            st.session_state.selected_workflow = None

elif st.session_state.selected_workflow == "data":
    with st.form("data_form", clear_on_submit=True):
        st.markdown("### 📊 Data Analysis Workflow")
        st.caption("Upload a zip dataset or enter a path to existing data")
        uploaded_zip = st.file_uploader(
            "Upload ZIP dataset (optional)",
            type=["zip"],
            help="Upload a zip file containing your dataset. Extracted temporarily, deleted on reset.",
        )
        if st.session_state.get("uploaded_data_path"):
            st.info(f"📁 Using uploaded data: `{st.session_state.uploaded_data_path}`")
        # data_path_manual = st.text_input(
        #     "Or enter data path manually",
        #     placeholder="e.g. /path/to/data.csv or /path/to/data_dir",
        # )
        query = st.text_input("Query", placeholder="What would you like to analyze?")
        submitted = st.form_submit_button("Run Data Analysis", type="primary")
        if submitted and query:
            path_to_use = None
            if uploaded_zip:
                cleanup_uploaded_data()  # Remove previous upload before saving new one
                extracted = save_and_extract_upload(uploaded_zip)
                if extracted and extracted.exists():
                    # Use absolute path and verify it exists
                    extracted = extracted.resolve()
                    st.session_state.uploaded_data_path = str(extracted)
                    path_to_use = str(extracted)
                    st.success(f"✅ File uploaded and extracted to: {path_to_use}")
                else:
                    st.error(f"Failed to process uploaded zip file. Extracted path: {extracted}")
            # elif data_path_manual.strip():
            #     path_to_use = data_path_manual.strip()
            elif st.session_state.get("uploaded_data_path"):
                path = Path(st.session_state.uploaded_data_path).resolve()
                if path.exists():
                    path_to_use = str(path)
                else:
                    st.warning(f"Previously uploaded path no longer exists: {path}")
                    cleanup_uploaded_data()
            if path_to_use:
                # Verify path exists before creating workflow config
                verify_path = Path(path_to_use).resolve()
                if not verify_path.exists():
                    st.error(f"Path does not exist: {path_to_use}")
                else:
                    workflow_config = {"type": "data", "path": str(verify_path), "query": query}
                    st.session_state.selected_workflow = None
            else:
                st.error("Please upload a zip file or enter a data path.")

elif st.session_state.selected_workflow == "experiment":
    with st.form("experiment_form", clear_on_submit=True):
        st.markdown("### 🧪 Experiment Workflow")
        st.caption("Upload a zip containing data_analysis.md or enter path manually")
        uploaded_exp_zip = st.file_uploader(
            "Upload ZIP with data analysis (optional)",
            type=["zip"],
            key="exp_upload",
            help="Zip containing data_analysis.md. Extracted temporarily, deleted on reset.",
        )
        if st.session_state.get("uploaded_experiment_path"):
            st.info(f"📁 Using: `{st.session_state.uploaded_experiment_path}`")
        # data_path_manual = st.text_input(
        #     "Or enter data analysis path manually",
        #     placeholder="Path to data_analysis.md (optional)",
        # )
        query = st.text_input("Experiment Query", placeholder="Describe your experiment...")
        submitted = st.form_submit_button("Run Experiment", type="primary")
        if submitted and query:
            path_to_use = None
            if uploaded_exp_zip:
                prev = st.session_state.get("uploaded_experiment_path")
                if prev:
                    _rm_upload_root(Path(prev))
                    if "uploaded_experiment_path" in st.session_state:
                        del st.session_state.uploaded_experiment_path
                extracted = save_and_extract_upload(uploaded_exp_zip)
                if extracted and extracted.exists():
                    extracted = extracted.resolve()
                    analysis_file = find_data_analysis_file(extracted)
                    if analysis_file and analysis_file.exists():
                        analysis_file = analysis_file.resolve()
                        st.session_state.uploaded_experiment_path = str(analysis_file)
                        path_to_use = str(analysis_file)
                        st.success(f"✅ Found analysis file: {path_to_use}")
                    else:
                        st.error(
                            f"Zip must contain data_analysis.md or analysis.md. Searched in: {extracted}"
                        )
                else:
                    st.error(f"Failed to process uploaded zip file. Extracted path: {extracted}")
            # elif data_path_manual.strip():
            #     path_to_use = data_path_manual.strip()
            elif st.session_state.get("uploaded_experiment_path"):
                p = Path(st.session_state.uploaded_experiment_path).resolve()
                if p.exists():
                    path_to_use = str(p)
                else:
                    st.warning(f"Previously uploaded path no longer exists: {p}")
                    if "uploaded_experiment_path" in st.session_state:
                        del st.session_state.uploaded_experiment_path
            if path_to_use:
                workflow_config = {"type": "experiment", "query": query, "path": path_to_use}
                st.session_state.selected_workflow = None
            else:
                st.error("Please upload a zip with data_analysis.md or enter a data analysis path.")

elif st.session_state.selected_workflow == "full":
    with st.form("full_form", clear_on_submit=True):
        st.markdown("### 🍎 Full Workflow")
        topic = st.text_input("Research Topic", placeholder="Enter your research topic...")
        st.caption("Data (for Data Analysis): upload zip or enter path")
        uploaded_full_zip = st.file_uploader(
            "Upload ZIP dataset (optional)",
            type=["zip"],
            key="full_upload",
            help="Zip dataset for Data Analysis. Extracted temporarily, deleted on reset.",
        )
        if st.session_state.get("uploaded_full_data_path"):
            st.info(f"📁 Using: `{st.session_state.uploaded_full_data_path}`")
        # data_path_manual = st.text_input(
        #     "Or enter data path manually",
        #     placeholder="Path to data file/dir (optional)",
        # )
        run_data = st.checkbox("Run Data Analysis", value=False)
        run_exp = st.checkbox("Run Experiment", value=False)
        submitted = st.form_submit_button("Run Full Workflow", type="primary")
        if submitted and topic:
            data_path_to_use = None
            if run_data:
                if uploaded_full_zip:
                    prev = st.session_state.get("uploaded_full_data_path")
                    if prev:
                        _rm_upload_root(Path(prev))
                        if "uploaded_full_data_path" in st.session_state:
                            del st.session_state.uploaded_full_data_path
                    extracted = save_and_extract_upload(uploaded_full_zip)
                    if extracted and extracted.exists():
                        extracted = extracted.resolve()
                        st.session_state.uploaded_full_data_path = str(extracted)
                        data_path_to_use = str(extracted)
                        st.success(f"✅ File uploaded and extracted to: {data_path_to_use}")
                    else:
                        st.error(
                            f"Failed to process uploaded zip file. Extracted path: {extracted}"
                        )
                        data_path_to_use = None
                # elif data_path_manual.strip():
                #     data_path_to_use = data_path_manual.strip()
                elif st.session_state.get("uploaded_full_data_path"):
                    p = Path(st.session_state.uploaded_full_data_path).resolve()
                    if p.exists():
                        data_path_to_use = str(p)
                    else:
                        st.warning(f"Previously uploaded path no longer exists: {p}")
                        if "uploaded_full_data_path" in st.session_state:
                            del st.session_state.uploaded_full_data_path
                if not data_path_to_use:
                    st.error("Run Data Analysis requires uploading a zip or entering a data path.")
                    data_path_to_use = None
            if data_path_to_use is not None or not run_data:
                workflow_config = {
                    "type": "full",
                    "query": topic,
                    "data_path": data_path_to_use,
                    "run_data": run_data,
                    "run_exp": run_exp,
                }
                st.session_state.selected_workflow = None

if workflow_config:
    # Allocate memo_X/workspace for this run
    memo_dir, workspace_path = allocate_memo_workspace()
    st.session_state.workspace_path = workspace_path

    # Add user message to chat
    if workflow_config["type"] == "ideation":
        user_msg = f"Ideation: {workflow_config['query']}"
    elif workflow_config["type"] == "data":
        user_msg = f"Data Analysis: {workflow_config['path']} - {workflow_config['query']}"
    elif workflow_config["type"] == "experiment":
        user_msg = f"Experiment: {workflow_config['query']}"
        if workflow_config.get("path"):
            user_msg += f" (Data: {workflow_config['path']})"
    else:  # full
        user_msg = f"Full Workflow: {workflow_config['query']}"
        if workflow_config.get("data_path"):
            user_msg += f" (Data: {workflow_config['data_path']})"
        if workflow_config.get("run_data"):
            user_msg += " [Data Analysis]"
        if workflow_config.get("run_exp"):
            user_msg += " [Experiment]"

    st.session_state.messages.append({"role": "user", "content": user_msg})

    # Execute workflow
    with st.chat_message("assistant"):
        loading_placeholder = st.empty()
        with loading_placeholder.container():
            st.markdown("Processing your request...")
            with st.spinner(""):
                if workflow_config["type"] == "ideation":
                    resp, intermediate_state = run_ideation(workflow_config.get("query"))
                elif workflow_config["type"] == "data":
                    resp, intermediate_state = run_data(
                        workflow_config["path"], workflow_config["query"]
                    )
                elif workflow_config["type"] == "experiment":
                    resp, intermediate_state = run_experiment(
                        workflow_config["query"], workflow_config.get("path")
                    )
                elif workflow_config["type"] == "full":
                    resp, intermediate_state = run_full(workflow_config)
                else:
                    resp, intermediate_state = "Unknown workflow type", []

        loading_placeholder.empty()
        stream_markdown(resp)
        render_intermediate_state(intermediate_state)
        msg = {"role": "assistant", "content": resp}
        if intermediate_state:
            msg["intermediate_state"] = intermediate_state
        st.session_state.messages.append(msg)

        metadata = {
            "workflow_type": workflow_config["type"],
            "query": workflow_config.get("query"),
            "path": workflow_config.get("path"),
        }
        if workflow_config["type"] == "full":
            metadata.update(
                {
                    "data_path": workflow_config.get("data_path"),
                    "run_data": workflow_config.get("run_data"),
                    "run_exp": workflow_config.get("run_exp"),
                }
            )

        save_chat_history(
            st.session_state.messages,
            workflow_type=workflow_config["type"],
            metadata=metadata,
            memo_dir=memo_dir,
        )
        st.session_state.last_saved_memo = str(memo_dir)

    st.rerun()


def parse_command(prompt):
    prompt = prompt.strip()
    if prompt.startswith("/ideation"):
        p = prompt.split(maxsplit=1)
        return {"type": "ideation", "query": p[1] if len(p) > 1 else None}
    if prompt.startswith("/data"):
        p = prompt.split(maxsplit=2)
        if len(p) < 3:
            return {"type": "error", "msg": "Usage: /data <path> <query>"}
        return {"type": "data", "path": p[1], "query": p[2]}
    if prompt.startswith("/experiment"):
        p = prompt.split(maxsplit=1)
        if len(p) < 2:
            return {"type": "error", "msg": "Usage: /experiment <query> [data_path]"}
        r = p[1].split(maxsplit=1)
        return {"type": "experiment", "query": r[0], "path": r[1] if len(r) > 1 else None}
    if prompt.startswith("/full"):
        p = prompt.split()
        cfg = {
            "type": "full",
            "query": p[1] if len(p) > 1 else None,
            "data_path": None,
            "run_data": False,
            "run_exp": False,
        }
        i = 2
        while i < len(p):
            if p[i] == "--data" and i + 1 < len(p):
                cfg["data_path"] = p[i + 1]
                cfg["run_data"] = True
                i += 2
            elif p[i] == "--experiment":
                cfg["run_exp"] = True
                i += 1
            else:
                i += 1
        return cfg
    return {"type": "ideation", "query": prompt}


# Chat input for general questions (use workflow buttons above for structured workflows)
if prompt := st.chat_input("Ask a question or select a workflow above"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        loading_placeholder = st.empty()
        with loading_placeholder.container():
            st.markdown("Processing your request...")
            with st.spinner(""):
                resp, intermediate_state = run_ideation(prompt)

        loading_placeholder.empty()
        stream_markdown(resp)
        render_intermediate_state(intermediate_state)
        msg = {"role": "assistant", "content": resp}
        if intermediate_state:
            msg["intermediate_state"] = intermediate_state
        st.session_state.messages.append(msg)

        memo_dir = save_chat_history(
            st.session_state.messages, workflow_type="ideation", metadata={"query": prompt}
        )
        st.session_state.last_saved_memo = str(memo_dir)

    st.rerun()