Agent Deep Research

version: 2
updates:
  # GitHub Actions
  - package-ecosystem: "github-actions"
    directory: "/"
    schedule:
      interval: "weekly"
      day: "monday"
    open-pull-requests-limit: 5
    commit-message:
      prefix: "chore(ci)"

blank_issues_enabled: false
contact_links:
  - name: Questions & Discussion
    url: https://github.com/24601/agent-deep-research/discussions
    about: Ask questions and discuss ideas in GitHub Discussions

name: CI

on:
  push:
    branches: [main]
  pull_request:
    branches: [main]

jobs:
  skill-validation:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v6

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.12'

      - name: Install uv
        uses: astral-sh/setup-uv@v6

      - name: Verify SKILL.md exists with valid frontmatter
        run: |
          test -f SKILL.md || { echo "SKILL.md not found"; exit 1; }
          head -10 SKILL.md | grep -q '^name:' || { echo "Missing 'name' in SKILL.md frontmatter"; exit 1; }
          head -10 SKILL.md | grep -q '^description:' || { echo "Missing 'description' in SKILL.md frontmatter"; exit 1; }
          echo "SKILL.md frontmatter is valid"

      - name: Check Python script syntax
        run: python3 -m py_compile scripts/research.py scripts/store.py scripts/upload.py scripts/state.py

      - name: Smoke test state.py
        run: uv run scripts/state.py --help

      - name: Verify reference files exist
        run: |
          test -f references/online_docs.md || { echo "Missing references/online_docs.md"; exit 1; }
          test -f references/file_search_guide.md || { echo "Missing references/file_search_guide.md"; exit 1; }
          echo "All reference files present"

      - name: Verify community files exist
        run: |
          test -f CHANGELOG.md || { echo "Missing CHANGELOG.md"; exit 1; }
          test -f CONTRIBUTING.md || { echo "Missing CONTRIBUTING.md"; exit 1; }
          echo "All community files present"

name: Release

on:
  push:
    tags:
      - 'v*'

permissions:
  contents: write

jobs:
  release:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v6

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.12'

      - name: Install uv
        uses: astral-sh/setup-uv@v6

      - name: Validate SKILL.md
        run: |
          test -f SKILL.md || { echo "SKILL.md not found"; exit 1; }
          head -10 SKILL.md | grep -q '^name:' || { echo "Missing 'name' in SKILL.md frontmatter"; exit 1; }
          head -10 SKILL.md | grep -q '^description:' || { echo "Missing 'description' in SKILL.md frontmatter"; exit 1; }
          echo "SKILL.md frontmatter is valid"

      - name: Check Python script syntax
        run: python3 -m py_compile scripts/research.py scripts/store.py scripts/upload.py scripts/state.py

      - name: Smoke test
        run: uv run scripts/state.py --help

      - name: Create GitHub Release
        uses: softprops/action-gh-release@v2
        with:
          generate_release_notes: true

# Python
__pycache__/
*.pyc
*.pyo
*.pyd
.Python

# OS
.DS_Store
Thumbs.db

# IDEs
.vscode/
.idea/

# Environment variables
.env
.env.local
.env.*.local

# Gemini Research local persistence
.gemini-research.json

# Test output artifacts
test-outputs/

# Letta conversation cache
.letta/

# Midscene logs
midscene_run/

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "google-genai>=1.0.0",
#     "rich>=13.0.0",
# ]
# ///
"""Start, monitor, and save Gemini Deep Research interactions.

Wraps the Gemini Interactions API to launch background deep-research
tasks, poll their status, and export the final report as Markdown.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path

from google import genai
from google.genai import types
from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.panel import Panel
from rich.spinner import Spinner
from rich.table import Table
from rich.text import Text

console = Console(stderr=True)

DEFAULT_AGENT = os.environ.get(
    "GEMINI_DEEP_RESEARCH_AGENT",
    "deep-research-pro-preview-12-2025",
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def get_api_key() -> str:
    """Resolve the API key from environment variables."""
    for var in ("GEMINI_DEEP_RESEARCH_API_KEY", "GOOGLE_API_KEY", "GEMINI_API_KEY"):
        key = os.environ.get(var)
        if key:
            return key
    console.print("[red]Error:[/red] No API key found.")
    console.print("Set one of: GEMINI_DEEP_RESEARCH_API_KEY, GOOGLE_API_KEY, GEMINI_API_KEY")
    sys.exit(1)


def get_client() -> genai.Client:
    """Create an authenticated GenAI client."""
    return genai.Client(api_key=get_api_key())


def get_state_path() -> Path:
    return Path(".gemini-research.json")


def load_state() -> dict:
    path = get_state_path()
    if not path.exists():
        return {"researchIds": [], "fileSearchStores": {}, "uploadOperations": {}}
    try:
        return json.loads(path.read_text())
    except (json.JSONDecodeError, OSError):
        return {"researchIds": [], "fileSearchStores": {}, "uploadOperations": {}}


def save_state(state: dict) -> None:
    get_state_path().write_text(json.dumps(state, indent=2) + "\n")


def add_research_id(interaction_id: str) -> None:
    """Track a research interaction ID in workspace state."""
    state = load_state()
    ids = state.setdefault("researchIds", [])
    if interaction_id not in ids:
        ids.append(interaction_id)
        save_state(state)


def record_research_completion(
    interaction_id: str, duration: int, grounded: bool,
) -> None:
    """Record a completed research run for adaptive polling."""
    state = load_state()
    history = state.setdefault("researchHistory", [])
    history.append({
        "id": interaction_id,
        "duration_seconds": duration,
        "grounded": grounded,
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    })
    # Keep last 50 entries to prevent unbounded growth
    state["researchHistory"] = history[-50:]
    save_state(state)


def _percentile(sorted_values: list[float], p: float) -> float:
    """Compute the p-th percentile (0-100) of a sorted list of values."""
    if not sorted_values:
        return 0.0
    k = (len(sorted_values) - 1) * (p / 100.0)
    f = int(k)
    c = f + 1
    if c >= len(sorted_values):
        return sorted_values[-1]
    return sorted_values[f] + (k - f) * (sorted_values[c] - sorted_values[f])


def _get_adaptive_poll_interval(
    elapsed: float, history: list[dict], grounded: bool,
) -> float:
    """Return poll interval based on historical completion times.

    Adapts the polling frequency so that we poll most aggressively during the
    window where research is most likely to finish (p25-p75 of past durations).
    Falls back to the fixed curve when insufficient history exists (<3 points).
    """
    # Filter history by grounded / non-grounded
    durations = sorted(
        entry["duration_seconds"]
        for entry in history
        if entry.get("grounded", False) == grounded
        and isinstance(entry.get("duration_seconds"), (int, float))
    )

    # Need at least 3 data points to build a meaningful distribution
    if len(durations) < 3:
        return _get_poll_interval(elapsed)

    min_d = durations[0]
    p25 = _percentile(durations, 25)
    p75 = _percentile(durations, 75)
    max_d = durations[-1]

    if elapsed < min_d:
        # Nothing ever finishes this fast -- poll slowly
        interval = 30.0
    elif elapsed < p25:
        # Some finish here -- moderate polling
        interval = 15.0
    elif elapsed <= p75:
        # Most likely completion window -- aggressive polling
        interval = 5.0
    elif elapsed <= max_d:
        # Tail end -- moderate
        interval = 15.0
    elif elapsed <= max_d * 1.5:
        # Past longest ever but within 1.5x -- slow down
        interval = 30.0
    else:
        # Unusually long -- very slow
        interval = 60.0

    # Clamp to [2, 120] seconds as fail-safe
    return max(2.0, min(120.0, interval))


def _write_output_dir(
    output_dir: str,
    interaction_id: str,
    interaction: object,
    report_text: str,
    duration_seconds: int | None = None,
) -> dict:
    """Write research results to a structured directory and return a compact summary."""
    base = Path(output_dir)
    research_dir = base / f"research-{interaction_id[:12]}"
    research_dir.mkdir(parents=True, exist_ok=True)

    # Write report.md
    report_path = research_dir / "report.md"
    report_path.write_text(report_text)

    # Build interaction data
    outputs_data = []
    sources: list[str] = []
    if interaction.outputs:
        for i, output in enumerate(interaction.outputs):
            text = getattr(output, "text", None)
            entry: dict = {"index": i, "text": text}
            outputs_data.append(entry)
            # Try to extract URLs from the text as sources
            if text:
                import re
                urls = re.findall(r'https?://[^\s\)>\]"\']+', text)
                sources.extend(urls)

    # Write interaction.json
    interaction_data = {
        "id": interaction_id,
        "status": getattr(interaction, "status", "unknown"),
        "outputCount": len(outputs_data),
        "outputs": outputs_data,
    }
    (research_dir / "interaction.json").write_text(
        json.dumps(interaction_data, indent=2, default=str) + "\n"
    )

    # Write sources.json (deduplicated)
    seen: set[str] = set()
    unique_sources: list[str] = []
    for url in sources:
        if url not in seen:
            seen.add(url)
            unique_sources.append(url)
    (research_dir / "sources.json").write_text(
        json.dumps(unique_sources, indent=2) + "\n"
    )

    # Write metadata.json
    metadata = {
        "id": interaction_id,
        "status": getattr(interaction, "status", "unknown"),
        "report_file": str(report_path),
        "report_size_bytes": len(report_text.encode("utf-8")),
        "output_count": len(outputs_data),
        "source_count": len(unique_sources),
    }
    if duration_seconds is not None:
        metadata["duration_seconds"] = duration_seconds
    (research_dir / "metadata.json").write_text(
        json.dumps(metadata, indent=2) + "\n"
    )

    # Build compact stdout summary (< 500 chars)
    summary_text = report_text[:200].replace("\n", " ").strip()
    if len(report_text) > 200:
        summary_text += "..."
    compact = {
        "id": interaction_id,
        "status": getattr(interaction, "status", "unknown"),
        "output_dir": str(research_dir),
        "report_file": str(report_path),
        "report_size_bytes": len(report_text.encode("utf-8")),
        "summary": summary_text,
    }
    if duration_seconds is not None:
        compact["duration_seconds"] = duration_seconds

    return compact


def resolve_store_name(name_or_alias: str) -> str:
    """Resolve a store display name to its resource name via state, or pass through."""
    if name_or_alias.startswith("fileSearchStores/"):
        return name_or_alias
    state = load_state()
    stores = state.get("fileSearchStores", {})
    if name_or_alias in stores:
        return stores[name_or_alias]
    return name_or_alias

# ---------------------------------------------------------------------------
# start subcommand
# ---------------------------------------------------------------------------

def cmd_start(args: argparse.Namespace) -> None:
    """Start a new deep research interaction."""
    client = get_client()
    query: str = args.query

    # Prepend report format if specified
    if args.report_format:
        format_map = {
            "executive_summary": "Executive Brief",
            "detailed_report": "Technical Deep Dive",
            "comprehensive": "Comprehensive Research Report",
        }
        label = format_map.get(args.report_format, args.report_format)
        query = f"[Report Format: {label}]\n\n{query}"

    # Handle follow-up: prepend context from previous interaction
    if args.follow_up:
        console.print(f"Loading previous research [bold]{args.follow_up}[/bold] for context...")
        try:
            prev = client.interactions.get(args.follow_up)
            if prev.outputs:
                prev_text = ""
                for output in prev.outputs:
                    text = getattr(output, "text", None)
                    if text:
                        prev_text = text  # use the last text output
                if prev_text:
                    query = (
                        f"[Follow-up to previous research]\n\n"
                        f"Previous findings:\n{prev_text[:4000]}\n\n"
                        f"New question:\n{query}"
                    )
        except Exception as exc:
            console.print(f"[yellow]Warning:[/yellow] Could not load previous research: {exc}")

    # Handle file attachment: upload to a temporary store
    file_search_store_names: list[str] | None = None

    if args.store:
        file_search_store_names = [resolve_store_name(args.store)]

    if args.file:
        filepath = Path(args.file).resolve()
        if not filepath.exists():
            console.print(f"[red]Error:[/red] File not found: {filepath}")
            sys.exit(1)
        if args.use_file_store:
            # Upload to a store for grounding
            console.print(f"Uploading [bold]{filepath.name}[/bold] to file search store...")
            store = client.file_search_stores.create(
                config={"display_name": f"research-{filepath.stem}"}
            )
            operation = client.file_search_stores.upload_to_file_search_store(
                file=str(filepath),
                file_search_store_name=store.name,
                config={"display_name": filepath.name},
            )
            while not operation.done:
                time.sleep(3)
                operation = client.operations.get(operation)
            console.print(f"[green]Uploaded to store:[/green] {store.name}")
            if file_search_store_names is None:
                file_search_store_names = []
            file_search_store_names.append(store.name)

            # Track in state
            st = load_state()
            st.setdefault("fileSearchStores", {})[f"research-{filepath.stem}"] = store.name
            save_state(st)
        else:
            # Inline file: append file contents to query (for smaller files)
            try:
                content = filepath.read_text(errors="replace")
                if len(content) > 100_000:
                    console.print(
                        "[yellow]Warning:[/yellow] File is large. "
                        "Consider using --use-file-store for better results."
                    )
                query = f"{query}\n\n---\nAttached file ({filepath.name}):\n{content}"
            except Exception as exc:
                console.print(f"[red]Error reading file:[/red] {exc}")
                sys.exit(1)

    # Build create kwargs
    create_kwargs: dict = {
        "input": query,
        "agent": DEFAULT_AGENT,
        "background": True,
    }
    if file_search_store_names:
        create_kwargs["config"] = {
            "file_search_store_names": file_search_store_names,
        }

    console.print("Starting deep research...")
    try:
        interaction = client.interactions.create(**create_kwargs)
    except Exception as exc:
        # Fallback: try without config if the SDK version doesn't support it
        if file_search_store_names and "config" in create_kwargs:
            console.print("[yellow]Note:[/yellow] Retrying without file search store config...")
            del create_kwargs["config"]
            try:
                interaction = client.interactions.create(**create_kwargs)
            except Exception as inner_exc:
                console.print(f"[red]Error:[/red] {inner_exc}")
                sys.exit(1)
        else:
            console.print(f"[red]Error:[/red] {exc}")
            sys.exit(1)

    interaction_id = interaction.id
    add_research_id(interaction_id)

    console.print(f"[green]Research started.[/green]")
    console.print(f"  ID: [bold]{interaction_id}[/bold]")
    console.print(f"  Status: {interaction.status}")
    console.print()
    console.print("Use [bold]research.py status[/bold] to check progress.")

    # If --output or --output-dir is set, poll until complete then save
    output_dir = getattr(args, "output_dir", None)
    grounded = file_search_store_names is not None
    adaptive_poll = not getattr(args, "no_adaptive_poll", False)
    if args.output or output_dir:
        _poll_and_save(
            client, interaction_id,
            output_path=args.output,
            output_dir=output_dir,
            show_thoughts=not args.no_thoughts,
            timeout=args.timeout,
            grounded=grounded,
            adaptive_poll=adaptive_poll,
        )
    else:
        # Print to stdout for machine consumption
        print(json.dumps({"id": interaction_id, "status": interaction.status}))


def _get_poll_interval(elapsed: float) -> float:
    """Return an adaptive poll interval based on elapsed time."""
    if elapsed < 30:
        return 5
    elif elapsed < 120:
        return 10
    elif elapsed < 600:
        return 30
    else:
        return 60


def _poll_and_save(
    client: genai.Client,
    interaction_id: str,
    output_path: str | None = None,
    output_dir: str | None = None,
    show_thoughts: bool = True,
    timeout: int = 1800,
    grounded: bool = False,
    adaptive_poll: bool = True,
) -> None:
    """Poll until research completes, then save the report."""
    console.print("Waiting for research to complete...")

    # Load history for adaptive polling
    history: list[dict] = []
    use_adaptive = False
    if adaptive_poll:
        try:
            state = load_state()
            history = state.get("researchHistory", [])
            # Need at least 3 matching entries to use adaptive
            matching = [
                e for e in history
                if e.get("grounded", False) == grounded
                and isinstance(e.get("duration_seconds"), (int, float))
            ]
            use_adaptive = len(matching) >= 3
        except Exception:
            pass  # Silently fall back to fixed curve

    if use_adaptive:
        console.print("[dim]Using adaptive polling (based on history).[/dim]")

    prev_output_count = 0
    start_time = time.monotonic()
    with Live(Spinner("dots", text="Researching..."), console=console, refresh_per_second=4) as live:
        while True:
            elapsed = time.monotonic() - start_time
            if elapsed > timeout:
                live.update(Text(f"Timed out after {int(elapsed)}s.", style="red bold"))
                console.print(f"[red]Error:[/red] Research timed out after {int(elapsed)} seconds.")
                console.print(f"Use [bold]research.py status {interaction_id}[/bold] to check later.")
                sys.exit(1)

            try:
                interaction = client.interactions.get(interaction_id)
            except Exception as exc:
                # Transient error -- log and retry
                interval = (
                    _get_adaptive_poll_interval(elapsed, history, grounded)
                    if use_adaptive
                    else _get_poll_interval(elapsed)
                )
                live.update(Text(f"Poll error (retrying): {exc}", style="yellow"))
                time.sleep(interval)
                continue

            status = interaction.status

            if show_thoughts and interaction.outputs:
                current_count = len(interaction.outputs)
                if current_count > prev_output_count:
                    # Show new thinking steps
                    for output in interaction.outputs[prev_output_count:]:
                        text = getattr(output, "text", None)
                        if text:
                            live.update(
                                Panel(
                                    Text(text[:500] + ("..." if len(text) > 500 else ""), style="dim"),
                                    title=f"Status: {status} ({int(elapsed)}s elapsed)",
                                    subtitle=f"Step {current_count}",
                                )
                            )
                    prev_output_count = current_count

            if status == "completed":
                live.update(Text("Research complete!", style="green bold"))
                break
            elif status in ("failed", "cancelled"):
                live.update(Text(f"Research {status}.", style="red bold"))
                console.print(f"[red]Research {status}.[/red]")
                sys.exit(1)

            interval = (
                _get_adaptive_poll_interval(elapsed, history, grounded)
                if use_adaptive
                else _get_poll_interval(elapsed)
            )
            time.sleep(interval)

    duration = int(time.monotonic() - start_time)

    # Record completion for future adaptive polling
    try:
        record_research_completion(interaction_id, duration, grounded)
    except Exception:
        pass  # Non-critical -- don't fail the save over history tracking

    # Extract final report
    report_text = ""
    if interaction.outputs:
        for output in reversed(interaction.outputs):
            text = getattr(output, "text", None)
            if text:
                report_text = text
                break

    if not report_text:
        console.print("[yellow]Warning:[/yellow] No text output found in completed research.")
        return

    # Write to output directory if specified
    if output_dir:
        compact = _write_output_dir(output_dir, interaction_id, interaction, report_text, duration)
        console.print()
        console.print(f"[green]Results saved to:[/green] {compact['output_dir']}")
        print(json.dumps(compact))
        return

    # Write to single file
    if output_path:
        Path(output_path).write_text(report_text)
        console.print()
        console.print(f"[green]Report saved to:[/green] {output_path}")

# ---------------------------------------------------------------------------
# status subcommand
# ---------------------------------------------------------------------------

def cmd_status(args: argparse.Namespace) -> None:
    """Check the status of a research interaction."""
    client = get_client()
    interaction_id: str = args.research_id

    try:
        interaction = client.interactions.get(interaction_id)
    except Exception as exc:
        console.print(f"[red]Error:[/red] {exc}")
        sys.exit(1)

    # Status summary
    status = interaction.status
    style = {"completed": "green", "failed": "red", "cancelled": "red"}.get(status, "yellow")
    console.print(f"Status: [{style}]{status}[/{style}]")
    console.print(f"ID: {interaction_id}")

    # Show outputs summary
    outputs = interaction.outputs or []
    if outputs:
        console.print(f"Outputs: {len(outputs)} step(s)")
        console.print()

        for i, output in enumerate(outputs):
            text = getattr(output, "text", None)
            if text:
                label = "Final Report" if i == len(outputs) - 1 and status == "completed" else f"Step {i + 1}"
                # Truncate for display
                preview = text[:300] + ("..." if len(text) > 300 else "")
                console.print(Panel(preview, title=label))
    else:
        console.print("[dim]No outputs yet.[/dim]")

    # Machine-readable on stdout
    result: dict = {"id": interaction_id, "status": status, "outputCount": len(outputs)}
    print(json.dumps(result))

# ---------------------------------------------------------------------------
# report subcommand
# ---------------------------------------------------------------------------

def cmd_report(args: argparse.Namespace) -> None:
    """Generate and save a markdown report from a completed interaction."""
    client = get_client()
    interaction_id: str = args.research_id

    try:
        interaction = client.interactions.get(interaction_id)
    except Exception as exc:
        console.print(f"[red]Error:[/red] {exc}")
        sys.exit(1)

    if interaction.status != "completed":
        console.print(
            f"[red]Error:[/red] Interaction is not completed. "
            f"Current status: {interaction.status}"
        )
        sys.exit(1)

    outputs = interaction.outputs or []
    if not outputs:
        console.print("[red]Error:[/red] No outputs found for this interaction.")
        sys.exit(1)

    # Build markdown report from outputs
    sections: list[str] = []
    sections.append(f"# Deep Research Report\n")
    sections.append(f"**Interaction ID:** `{interaction_id}`\n")
    sections.append(f"**Status:** {interaction.status}\n")
    sections.append("---\n")

    for i, output in enumerate(outputs):
        text = getattr(output, "text", None)
        if text:
            if i == len(outputs) - 1:
                sections.append(text)
            else:
                sections.append(f"### Research Step {i + 1}\n")
                sections.append(text)
                sections.append("\n---\n")

    report = "\n".join(sections)

    output_dir = getattr(args, "output_dir", None)
    if output_dir:
        compact = _write_output_dir(output_dir, interaction_id, interaction, report)
        console.print(f"[green]Results saved to:[/green] {compact['output_dir']}")
        print(json.dumps(compact))
        return

    output_path = args.output or f"research-report-{interaction_id[:8]}.md"
    Path(output_path).write_text(report)
    console.print(f"[green]Report saved to:[/green] {output_path}")

# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="research",
        description="Gemini Deep Research: start, monitor, and save research interactions",
    )
    sub = parser.add_subparsers(dest="command")

    # start (default)
    start_p = sub.add_parser("start", help="Start a new deep research interaction (default)")
    start_p.add_argument("query", help="The research query or instructions")
    start_p.add_argument(
        "--file", metavar="PATH",
        help="Attach a file to the research (inlined or uploaded to store)",
    )
    start_p.add_argument(
        "--use-file-store", action="store_true",
        help="Upload attached file to a file search store for grounding",
    )
    start_p.add_argument(
        "--store", metavar="NAME",
        help="Use a pre-existing file search store for grounding (name or resource ID)",
    )
    start_p.add_argument(
        "--report-format",
        choices=["executive_summary", "detailed_report", "comprehensive"],
        help="Desired report format",
    )
    start_p.add_argument(
        "--follow-up", metavar="ID",
        help="Continue from a previous research interaction",
    )
    start_p.add_argument(
        "--output", "-o", metavar="PATH",
        help="Wait for completion and save report to this path",
    )
    start_p.add_argument(
        "--no-thoughts", action="store_true",
        help="Suppress thinking step display during polling",
    )
    start_p.add_argument(
        "--timeout", type=int, default=1800,
        help="Maximum seconds to wait when --output is used (default: 1800)",
    )
    start_p.add_argument(
        "--output-dir", metavar="DIR",
        help="Wait for completion and save structured results to this directory",
    )
    start_p.add_argument(
        "--no-adaptive-poll", action="store_true",
        help="Disable history-adaptive polling; use fixed interval curve instead",
    )

    # status
    status_p = sub.add_parser("status", help="Check research interaction status")
    status_p.add_argument("research_id", help="The interaction ID")

    # report
    report_p = sub.add_parser("report", help="Save a markdown report from completed research")
    report_p.add_argument("research_id", help="The interaction ID")
    report_p.add_argument("--output", "-o", metavar="PATH", help="Output file path")
    report_p.add_argument(
        "--output-dir", metavar="DIR",
        help="Save structured results to this directory",
    )

    return parser


def main(argv: list[str] | None = None) -> None:
    parser = build_parser()
    args = parser.parse_args(argv)

    commands = {
        "start": cmd_start,
        "status": cmd_status,
        "report": cmd_report,
    }

    if args.command is None:
        # Default to start if a bare query is provided
        # Re-parse with start as default
        if argv is None:
            argv = sys.argv[1:]
        if argv and not argv[0].startswith("-") and argv[0] not in commands:
            argv = ["start"] + list(argv)
            args = parser.parse_args(argv)

    handler = commands.get(args.command)
    if handler is None:
        parser.print_help()
        sys.exit(1)
    handler(args)


if __name__ == "__main__":
    main()

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "rich>=13.0.0",
# ]
# ///
"""Manage workspace state for Gemini Deep Research.

Reads and manages .gemini-research.json which tracks research IDs,
file search store mappings, and upload operations.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from rich.console import Console
from rich.table import Table

console = Console(stderr=True)

# ---------------------------------------------------------------------------
# State helpers
# ---------------------------------------------------------------------------

def get_state_path() -> Path:
    """Return the path to the workspace state file."""
    return Path(".gemini-research.json")


def load_state() -> dict:
    """Load workspace state from disk, returning empty defaults if missing."""
    path = get_state_path()
    if not path.exists():
        return {"researchIds": [], "fileSearchStores": {}, "uploadOperations": {}}
    try:
        return json.loads(path.read_text())
    except (json.JSONDecodeError, OSError) as exc:
        console.print(f"[yellow]Warning:[/yellow] failed to read state file: {exc}")
        return {"researchIds": [], "fileSearchStores": {}, "uploadOperations": {}}


def save_state(state: dict) -> None:
    """Persist workspace state to disk."""
    get_state_path().write_text(json.dumps(state, indent=2) + "\n")

# ---------------------------------------------------------------------------
# Subcommands
# ---------------------------------------------------------------------------

def cmd_show(_args: argparse.Namespace) -> None:
    """Display full workspace state."""
    state = load_state()
    use_json = getattr(_args, "json", False)

    if use_json:
        # Emit full state (excluding internal caches) as JSON to stdout
        output = {
            "researchIds": state.get("researchIds", []),
            "fileSearchStores": state.get("fileSearchStores", {}),
            "uploadOperations": state.get("uploadOperations", {}),
        }
        print(json.dumps(output, indent=2))
        return

    if not any(state.get(k) for k in ("researchIds", "fileSearchStores", "uploadOperations")):
        console.print("[dim]No workspace state found.[/dim]")
        return

    # Research IDs
    ids = state.get("researchIds", [])
    if ids:
        table = Table(title="Research Interactions")
        table.add_column("#", style="dim")
        table.add_column("Interaction ID")
        for i, rid in enumerate(ids, 1):
            table.add_row(str(i), rid)
        console.print(table)
    else:
        console.print("[dim]No research interactions tracked.[/dim]")

    console.print()

    # File search stores
    stores = state.get("fileSearchStores", {})
    if stores:
        table = Table(title="File Search Stores")
        table.add_column("Display Name")
        table.add_column("Resource Name")
        for name, resource in stores.items():
            table.add_row(name, resource)
        console.print(table)
    else:
        console.print("[dim]No file search stores tracked.[/dim]")

    console.print()

    # Upload operations
    ops = state.get("uploadOperations", {})
    if ops:
        table = Table(title="Upload Operations")
        table.add_column("ID", style="dim")
        table.add_column("Status")
        table.add_column("Path")
        table.add_column("Store")
        table.add_column("Progress")
        for op_id, op in ops.items():
            total = op.get("totalFiles", 0)
            done = op.get("completedFiles", 0) + op.get("skippedFiles", 0)
            pct = f"{round(done / total * 100)}%" if total else "N/A"
            status = op.get("status", "unknown")
            style = {"completed": "green", "failed": "red", "in_progress": "yellow"}.get(status, "")
            table.add_row(
                op_id[:12],
                f"[{style}]{status}[/{style}]" if style else status,
                op.get("path", ""),
                op.get("storeName", ""),
                pct,
            )
        console.print(table)
    else:
        console.print("[dim]No upload operations tracked.[/dim]")


def cmd_research(_args: argparse.Namespace) -> None:
    """List tracked research IDs."""
    state = load_state()
    ids = state.get("researchIds", [])
    use_json = getattr(_args, "json", False)

    if use_json:
        print(json.dumps(ids))
        return

    if not ids:
        console.print("[dim]No research interactions tracked.[/dim]")
        return
    table = Table(title="Research Interactions")
    table.add_column("#", style="dim")
    table.add_column("Interaction ID")
    for i, rid in enumerate(ids, 1):
        table.add_row(str(i), rid)
    console.print(table)


def cmd_stores(_args: argparse.Namespace) -> None:
    """List tracked store mappings."""
    state = load_state()
    stores = state.get("fileSearchStores", {})
    use_json = getattr(_args, "json", False)

    if use_json:
        result = [{"displayName": k, "name": v} for k, v in stores.items()]
        print(json.dumps(result))
        return

    if not stores:
        console.print("[dim]No file search stores tracked.[/dim]")
        return
    table = Table(title="File Search Stores")
    table.add_column("Display Name")
    table.add_column("Resource Name")
    for name, resource in stores.items():
        table.add_row(name, resource)
    console.print(table)


def cmd_clear(_args: argparse.Namespace) -> None:
    """Reset workspace state."""
    path = get_state_path()
    if not path.exists():
        console.print("[dim]No state file to clear.[/dim]")
        return

    if not _args.yes:
        if not sys.stdin.isatty():
            # Non-interactive context (e.g. AI agent): auto-accept
            pass
        else:
            console.print(f"This will delete [bold]{path}[/bold]. Use -y to skip this prompt.")
            try:
                answer = input("Continue? [y/N] ").strip().lower()
            except (EOFError, KeyboardInterrupt):
                print()
                sys.exit(1)
            if answer not in ("y", "yes"):
                console.print("[dim]Aborted.[/dim]")
                return

    path.unlink()
    console.print("[green]Workspace state cleared.[/green]")

# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="state",
        description="Manage Gemini Deep Research workspace state (.gemini-research.json)",
    )
    parser.add_argument(
        "--json", action="store_true", dest="json",
        help="Output JSON to stdout for programmatic consumption",
    )
    sub = parser.add_subparsers(dest="command")

    sub.add_parser("show", help="Display full workspace state (default)")
    sub.add_parser("research", help="List tracked research interaction IDs")
    sub.add_parser("stores", help="List tracked file search store mappings")

    clear_p = sub.add_parser("clear", help="Reset workspace state")
    clear_p.add_argument("-y", "--yes", action="store_true", help="Skip confirmation prompt")

    return parser


def main(argv: list[str] | None = None) -> None:
    parser = build_parser()
    args = parser.parse_args(argv)

    commands = {
        "show": cmd_show,
        "research": cmd_research,
        "stores": cmd_stores,
        "clear": cmd_clear,
        None: cmd_show,  # default
    }

    handler = commands.get(args.command, cmd_show)
    handler(args)


if __name__ == "__main__":
    main()

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "google-genai>=1.0.0",
#     "rich>=13.0.0",
# ]
# ///
"""Manage Gemini File Search stores (corpora).

Provides create, list, delete, and query operations for file search
stores used with Gemini RAG grounding.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

from google import genai
from google.genai import types
from rich.console import Console
from rich.table import Table

console = Console(stderr=True)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def get_api_key() -> str:
    """Resolve the API key from environment variables."""
    for var in ("GEMINI_DEEP_RESEARCH_API_KEY", "GOOGLE_API_KEY", "GEMINI_API_KEY"):
        key = os.environ.get(var)
        if key:
            return key
    console.print("[red]Error:[/red] No API key found.")
    console.print("Set one of: GEMINI_DEEP_RESEARCH_API_KEY, GOOGLE_API_KEY, GEMINI_API_KEY")
    sys.exit(1)


def get_client() -> genai.Client:
    """Create an authenticated GenAI client."""
    return genai.Client(api_key=get_api_key())


def get_state_path() -> Path:
    """Return the path to the workspace state file."""
    return Path(".gemini-research.json")


def load_state() -> dict:
    path = get_state_path()
    if not path.exists():
        return {"researchIds": [], "fileSearchStores": {}, "uploadOperations": {}}
    try:
        return json.loads(path.read_text())
    except (json.JSONDecodeError, OSError):
        return {"researchIds": [], "fileSearchStores": {}, "uploadOperations": {}}


def save_state(state: dict) -> None:
    get_state_path().write_text(json.dumps(state, indent=2) + "\n")


def get_default_model() -> str:
    """Return the model to use for file search queries."""
    return os.environ.get(
        "GEMINI_DEEP_RESEARCH_MODEL",
        os.environ.get("GEMINI_MODEL", "models/gemini-flash-latest"),
    )

# ---------------------------------------------------------------------------
# Subcommands
# ---------------------------------------------------------------------------

def cmd_create(args: argparse.Namespace) -> None:
    """Create a new file search store."""
    client = get_client()
    display_name: str = args.name

    console.print(f"Creating store [bold]{display_name}[/bold]...")
    store = client.file_search_stores.create(config={"display_name": display_name})

    # Persist mapping
    state = load_state()
    state.setdefault("fileSearchStores", {})[display_name] = store.name
    save_state(state)

    console.print(f"[green]Created store:[/green] {store.name} ({display_name})")
    # Also emit machine-readable output to stdout
    print(json.dumps({"name": store.name, "displayName": display_name}))


def cmd_list(_args: argparse.Namespace) -> None:
    """List all file search stores."""
    client = get_client()

    stores: list[dict] = []
    for store in client.file_search_stores.list():
        display = getattr(store, "display_name", None)
        if display is None:
            cfg = getattr(store, "config", None)
            display = getattr(cfg, "display_name", "") if cfg else ""
        stores.append({"name": store.name, "displayName": display})

    if not stores:
        console.print("[dim]No file search stores found.[/dim]")
        print(json.dumps([]))
        return

    table = Table(title="File Search Stores")
    table.add_column("Resource Name")
    table.add_column("Display Name")
    for s in stores:
        table.add_row(s["name"], s["displayName"])
    console.print(table)

    # Machine-readable on stdout
    print(json.dumps(stores, indent=2))


def cmd_delete(args: argparse.Namespace) -> None:
    """Delete a file search store."""
    client = get_client()
    store_id: str = args.id
    force: bool = args.force

    if not force:
        if not sys.stdin.isatty():
            # Non-interactive context (e.g. AI agent): auto-accept
            force = True
        else:
            console.print(f"Deleting store [bold]{store_id}[/bold]. Use --force to skip this prompt.")
            try:
                answer = input("Continue? [y/N] ").strip().lower()
            except (EOFError, KeyboardInterrupt):
                print()
                sys.exit(1)
            if answer not in ("y", "yes"):
                console.print("[dim]Aborted.[/dim]")
                return

    console.print(f"Deleting store [bold]{store_id}[/bold]...")
    client.file_search_stores.delete(name=store_id, config={"force": force})
    console.print(f"[green]Deleted store:[/green] {store_id}")

    # Clean up local state: remove the store mapping
    state = load_state()
    stores = state.get("fileSearchStores", {})
    removed = [k for k, v in stores.items() if v == store_id or k == store_id]
    for k in removed:
        del stores[k]
    if removed:
        save_state(state)
        console.print(f"[dim]Removed {len(removed)} local mapping(s).[/dim]")


def cmd_query(args: argparse.Namespace) -> None:
    """Query a file search store with grounded generation."""
    import time as _time

    client = get_client()
    store_name: str = args.id
    question: str = args.question
    model = get_default_model()
    output_dir = getattr(args, "output_dir", None)

    console.print(f"Querying store [bold]{store_name}[/bold]...")
    start = _time.monotonic()
    try:
        response = client.models.generate_content(
            model=model,
            contents=question,
            config=types.GenerateContentConfig(
                tools=[
                    types.Tool(
                        file_search=types.FileSearch(
                            file_search_store_names=[store_name],
                        )
                    )
                ]
            ),
        )
        text = response.text if response.text else "No response generated."
    except Exception as exc:
        console.print(f"[red]Query failed:[/red] {exc}")
        sys.exit(1)
    duration = int(_time.monotonic() - start)

    if output_dir:
        import re
        base = Path(output_dir)
        ts = _time.strftime("%Y%m%d-%H%M%S")
        query_dir = base / f"query-{ts}"
        query_dir.mkdir(parents=True, exist_ok=True)

        (query_dir / "response.md").write_text(text)
        metadata = {
            "store": store_name,
            "question": question,
            "model": model,
            "response_size_bytes": len(text.encode("utf-8")),
            "duration_seconds": duration,
        }
        (query_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")

        summary_text = text[:200].replace("\n", " ").strip()
        if len(text) > 200:
            summary_text += "..."
        compact = {
            "output_dir": str(query_dir),
            "response_file": str(query_dir / "response.md"),
            "response_size_bytes": len(text.encode("utf-8")),
            "duration_seconds": duration,
            "summary": summary_text,
        }
        console.print(f"[green]Results saved to:[/green] {query_dir}")
        print(json.dumps(compact))
        return

    # Output answer to stdout (rich formatting on stderr)
    console.print("[green]Answer:[/green]")
    print(text)

# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="store",
        description="Manage Gemini File Search stores",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    # create
    create_p = sub.add_parser("create", help="Create a new file search store")
    create_p.add_argument("name", help="Display name for the store")

    # list
    sub.add_parser("list", help="List all file search stores")

    # delete
    del_p = sub.add_parser("delete", help="Delete a file search store")
    del_p.add_argument("id", help="Resource name of the store (e.g. fileSearchStores/...)")
    del_p.add_argument("--force", action="store_true", help="Force delete even if store contains documents")

    # query
    query_p = sub.add_parser("query", help="Query a file search store")
    query_p.add_argument("id", help="Resource name of the store")
    query_p.add_argument("question", help="The question to ask")
    query_p.add_argument(
        "--output-dir", metavar="DIR",
        help="Save response and metadata to this directory",
    )

    return parser


def main(argv: list[str] | None = None) -> None:
    parser = build_parser()
    args = parser.parse_args(argv)

    commands = {
        "create": cmd_create,
        "list": cmd_list,
        "delete": cmd_delete,
        "query": cmd_query,
    }

    handler = commands.get(args.command)
    if handler is None:
        parser.print_help()
        sys.exit(1)
    handler(args)


if __name__ == "__main__":
    main()

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "google-genai>=1.0.0",
#     "rich>=13.0.0",
# ]
# ///
"""Upload files to a Gemini File Search store.

Supports single files and recursive directory uploads with MIME type
validation, smart-sync (skip unchanged), and progress tracking.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import mimetypes
import os
import sys
import time
import uuid
from pathlib import Path

from google import genai
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn

console = Console(stderr=True)

# ---------------------------------------------------------------------------
# MIME type maps (derived from docs/file-search-mime-types.md)
# ---------------------------------------------------------------------------

# Tier 1: validated MIME types that work natively
VALIDATED_MIME: dict[str, str] = {
    ".pdf": "application/pdf",
    ".xml": "application/xml",
    ".txt": "text/plain",
    ".text": "text/plain",
    ".log": "text/plain",
    ".out": "text/plain",
    ".env": "text/plain",
    ".gitignore": "text/plain",
    ".gitattributes": "text/plain",
    ".dockerignore": "text/plain",
    ".html": "text/html",
    ".htm": "text/html",
    ".md": "text/markdown",
    ".markdown": "text/markdown",
    ".mdown": "text/markdown",
    ".mkd": "text/markdown",
    ".c": "text/x-c",
    ".h": "text/x-c",
    ".java": "text/x-java",
    ".kt": "text/x-kotlin",
    ".kts": "text/x-kotlin",
    ".go": "text/x-go",
    ".py": "text/x-python",
    ".pyw": "text/x-python",
    ".pyx": "text/x-python",
    ".pyi": "text/x-python",
    ".pl": "text/x-perl",
    ".pm": "text/x-perl",
    ".t": "text/x-perl",
    ".pod": "text/x-perl",
    ".lua": "text/x-lua",
    ".erl": "text/x-erlang",
    ".hrl": "text/x-erlang",
    ".tcl": "text/x-tcl",
    ".bib": "text/x-bibtex",
    ".diff": "text/x-diff",
}

# Tier 2: known text extensions that fall back to text/plain
TEXT_FALLBACK_EXTENSIONS: set[str] = {
    # JavaScript / TypeScript
    ".js", ".mjs", ".cjs", ".jsx",
    ".ts", ".mts", ".cts", ".tsx",
    ".json", ".jsonc", ".json5",
    # Web
    ".css", ".scss", ".sass", ".less", ".styl",
    ".vue", ".svelte", ".astro",
    # Shell / Scripting
    ".sh", ".bash", ".zsh", ".fish", ".ksh",
    ".bat", ".cmd", ".ps1", ".psm1",
    # Config
    ".yaml", ".yml", ".toml", ".ini", ".cfg", ".conf",
    ".properties", ".editorconfig", ".prettierrc",
    ".eslintrc", ".babelrc", ".npmrc",
    # Other languages
    ".rb", ".php", ".rs", ".swift", ".scala", ".clj",
    ".ex", ".exs", ".hs", ".ml", ".fs", ".fsx",
    ".r", ".jl", ".nim", ".zig", ".dart",
    ".coffee", ".elm", ".v", ".cr", ".groovy",
    ".gradle", ".cmake", ".makefile", ".mk",
    ".dockerfile", ".tf", ".hcl",
    ".sql", ".graphql", ".gql", ".proto",
    ".csv", ".tsv", ".rst", ".adoc", ".tex", ".latex",
    ".sbt", ".pom",
}

# Tier 3: binary extensions that must be rejected
BINARY_EXTENSIONS: set[str] = {
    ".exe", ".dll", ".so", ".dylib", ".a", ".lib",
    ".zip", ".tar", ".gz", ".bz2", ".7z", ".rar", ".xz",
    ".png", ".jpg", ".jpeg", ".gif", ".bmp", ".ico", ".svg", ".webp",
    ".mp3", ".mp4", ".wav", ".avi", ".mkv", ".mov", ".flac", ".ogg",
    ".class", ".pyc", ".pyo", ".o", ".obj",
    ".wasm", ".bin", ".dat",
    ".ttf", ".otf", ".woff", ".woff2", ".eot",
}

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def get_api_key() -> str:
    for var in ("GEMINI_DEEP_RESEARCH_API_KEY", "GOOGLE_API_KEY", "GEMINI_API_KEY"):
        key = os.environ.get(var)
        if key:
            return key
    console.print("[red]Error:[/red] No API key found.")
    console.print("Set one of: GEMINI_DEEP_RESEARCH_API_KEY, GOOGLE_API_KEY, GEMINI_API_KEY")
    sys.exit(1)


def get_client() -> genai.Client:
    return genai.Client(api_key=get_api_key())


def get_state_path() -> Path:
    return Path(".gemini-research.json")


def load_state() -> dict:
    path = get_state_path()
    if not path.exists():
        return {"researchIds": [], "fileSearchStores": {}, "uploadOperations": {}}
    try:
        return json.loads(path.read_text())
    except (json.JSONDecodeError, OSError):
        return {"researchIds": [], "fileSearchStores": {}, "uploadOperations": {}}


def save_state(state: dict) -> None:
    get_state_path().write_text(json.dumps(state, indent=2) + "\n")


def resolve_mime(filepath: Path) -> str | None:
    """Return MIME type for a file, or None if unsupported.

    Tier 1: validated native types.
    Tier 2: known text files -> text/plain fallback.
    Tier 3: binary -> None (rejected).
    """
    ext = filepath.suffix.lower()
    # Check special dotfiles (no suffix but known names)
    name_lower = filepath.name.lower()
    if name_lower in (".gitignore", ".gitattributes", ".dockerignore",
                       ".editorconfig", ".prettierrc", ".eslintrc",
                       ".babelrc", ".npmrc", ".env"):
        return VALIDATED_MIME.get(name_lower, "text/plain")

    if ext in VALIDATED_MIME:
        return VALIDATED_MIME[ext]
    if ext in TEXT_FALLBACK_EXTENSIONS:
        return "text/plain"
    if ext in BINARY_EXTENSIONS:
        return None

    # Unknown extension: try system mimetypes, accept text/* only
    guessed, _ = mimetypes.guess_type(str(filepath))
    if guessed and guessed.startswith("text/"):
        return "text/plain"

    # Default: reject unknown
    return None


def file_hash(filepath: Path) -> str:
    """Compute SHA-256 hash of a file for smart-sync."""
    h = hashlib.sha256()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def collect_files(
    root: Path,
    extensions: set[str] | None = None,
) -> list[Path]:
    """Recursively collect uploadable files from a directory."""
    files: list[Path] = []
    for p in sorted(root.rglob("*")):
        if not p.is_file():
            continue
        if extensions and p.suffix.lower() not in extensions:
            continue
        if resolve_mime(p) is not None:
            files.append(p)
    return files


def load_hash_cache(state: dict, store_name: str) -> dict[str, str]:
    """Load the per-store file hash cache from state."""
    return state.get("_hashCache", {}).get(store_name, {})


def save_hash_cache(state: dict, store_name: str, cache: dict[str, str]) -> None:
    state.setdefault("_hashCache", {})[store_name] = cache
    save_state(state)

# ---------------------------------------------------------------------------
# Upload logic
# ---------------------------------------------------------------------------

def upload_files(
    client: genai.Client,
    files: list[Path],
    store_name: str,
    smart_sync: bool = False,
) -> dict:
    """Upload a list of files to a store, returning an operation summary."""
    state = load_state()
    hash_cache = load_hash_cache(state, store_name)

    completed = 0
    skipped = 0
    failed = 0
    failed_list: list[dict] = []

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        console=console,
    ) as progress:
        task = progress.add_task("Uploading...", total=len(files))

        for filepath in files:
            rel = str(filepath)
            mime = resolve_mime(filepath)
            if mime is None:
                failed += 1
                failed_list.append({"file": rel, "error": "Unsupported file type"})
                progress.advance(task)
                continue

            # Compute hash for smart-sync comparison and cache update
            current_hash = file_hash(filepath)

            # Smart-sync: skip if hash unchanged
            if smart_sync and hash_cache.get(rel) == current_hash:
                skipped += 1
                progress.update(task, description=f"Skipped: {filepath.name}")
                progress.advance(task)
                continue

            try:
                progress.update(task, description=f"Uploading: {filepath.name}")
                operation = client.file_search_stores.upload_to_file_search_store(
                    file=str(filepath),
                    file_search_store_name=store_name,
                    config={"display_name": filepath.name},
                )
                # Poll until done
                while not operation.done:
                    time.sleep(2)
                    operation = client.operations.get(operation)

                completed += 1
                # Always update hash cache on successful upload (enables future smart-sync)
                hash_cache[rel] = current_hash
            except Exception as exc:
                failed += 1
                failed_list.append({"file": rel, "error": str(exc)})

            progress.advance(task)

    # Always persist hash cache so future --smart-sync runs can skip unchanged files
    save_hash_cache(state, store_name, hash_cache)

    return {
        "totalFiles": len(files),
        "completedFiles": completed,
        "skippedFiles": skipped,
        "failedFiles": failed,
        "failedFilesList": failed_list,
    }

# ---------------------------------------------------------------------------
# Subcommands
# ---------------------------------------------------------------------------

def cmd_upload(args: argparse.Namespace) -> None:
    """Upload files or directories to a file search store."""
    client = get_client()
    target = Path(args.path).resolve()
    store_name: str = args.store_name
    smart_sync: bool = args.smart_sync
    extensions: set[str] | None = None
    if args.extensions:
        # Accept both comma-separated and space-separated (via nargs)
        raw = args.extensions if isinstance(args.extensions, list) else [args.extensions]
        parts: list[str] = []
        for item in raw:
            parts.extend(item.replace(",", " ").split())
        extensions = {
            ext if ext.startswith(".") else f".{ext}"
            for ext in parts
            if ext.strip()
        }

    if not target.exists():
        console.print(f"[red]Error:[/red] Path not found: {target}")
        sys.exit(1)

    # Collect files
    if target.is_file():
        mime = resolve_mime(target)
        if mime is None:
            console.print(f"[red]Error:[/red] Unsupported file type: {target.suffix}")
            sys.exit(1)
        files = [target]
    elif target.is_dir():
        files = collect_files(target, extensions)
        if not files:
            console.print("[yellow]No uploadable files found.[/yellow]")
            sys.exit(0)
        console.print(f"Found [bold]{len(files)}[/bold] files to upload.")
    else:
        console.print(f"[red]Error:[/red] Path is not a file or directory: {target}")
        sys.exit(1)

    # Record operation in state
    op_id = str(uuid.uuid4())[:8]
    state = load_state()
    state.setdefault("uploadOperations", {})[op_id] = {
        "id": op_id,
        "status": "in_progress",
        "path": str(target),
        "storeName": store_name,
        "smartSync": smart_sync,
        "totalFiles": len(files),
        "completedFiles": 0,
        "skippedFiles": 0,
        "failedFiles": 0,
        "failedFilesList": [],
        "startedAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }
    save_state(state)

    console.print(f"Upload operation: [bold]{op_id}[/bold]")
    result = upload_files(client, files, store_name, smart_sync)

    # Update operation in state
    state = load_state()
    op = state["uploadOperations"][op_id]
    op.update(result)
    op["status"] = "failed" if result["failedFiles"] == result["totalFiles"] else "completed"
    op["completedAt"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    save_state(state)

    # Summary
    console.print()
    console.print(f"[green]Completed:[/green] {result['completedFiles']}")
    console.print(f"[yellow]Skipped:[/yellow]   {result['skippedFiles']}")
    console.print(f"[red]Failed:[/red]    {result['failedFiles']}")
    if result["failedFilesList"]:
        console.print("[red]Failed files:[/red]")
        for f in result["failedFilesList"]:
            console.print(f"  {f['file']}: {f['error']}")

    print(json.dumps({"operationId": op_id, **result}))


def cmd_status(args: argparse.Namespace) -> None:
    """Check upload operation status from local state."""
    state = load_state()
    ops = state.get("uploadOperations", {})
    op = ops.get(args.operation_id)
    if not op:
        console.print(f"[red]Error:[/red] Operation not found: {args.operation_id}")
        sys.exit(1)
    print(json.dumps(op, indent=2))

# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="upload",
        description="Upload files to a Gemini File Search store",
    )
    parser.add_argument("path", nargs="?", help="Path to file or directory to upload")
    parser.add_argument("store_name", nargs="?", help="Resource name of the file search store")
    parser.add_argument(
        "--smart-sync", action="store_true",
        help="Skip uploading files that have not changed (hash comparison)",
    )
    parser.add_argument(
        "--extensions", nargs="*",
        help="File extensions to include (comma or space separated, e.g. py,ts,md or .py .ts .md)",
    )
    parser.add_argument(
        "--status",
        dest="operation_id",
        help="Check status of an upload operation instead of uploading",
    )

    return parser


def main(argv: list[str] | None = None) -> None:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.operation_id:
        cmd_status(args)
        return

    if not args.path or not args.store_name:
        parser.error("path and store_name are required for upload (or use --status)")

    cmd_upload(args)


if __name__ == "__main__":
    main()