/* Copyright 2026 Marimo. All rights reserved. */

export function getAgentPrompt(filename: string) {
  return `
  I am currently editing a marimo notebook.
  You can read or write to the notebook at ${filename}

  If you make edits to the notebook, only edit the contents inside the function decorator with @app.cell.
  marimo will automatically handle adding the parameters and return statement of the function. For example,
  for each edit, just return:

  \`\`\`
  @app.cell
  def _():
    <your code here>
    return
  \`\`\`

  ## Marimo fundamentals

  Marimo is a reactive notebook that differs from traditional notebooks in key ways:

  - Cells execute automatically when their dependencies change
  - Variables cannot be redeclared across cells
  - The notebook forms a directed acyclic graph (DAG)
  - The last expression in a cell is automatically displayed
  - UI elements are reactive and update the notebook automatically

  ## Code Requirements

  1. All code must be complete and runnable
  2. Follow consistent coding style throughout
  3. Include descriptive variable names and helpful comments
  4. Import all modules in the first cell, always including \`import marimo as mo\`
  5. Never redeclare variables across cells
  6. Ensure no cycles in notebook dependency graph
  7. The last expression in a cell is automatically displayed, just like in Jupyter notebooks.
  8. Don't include comments in markdown cells
  9. Don't include comments in SQL cells
  10. Never define anything using \`global\`.

  ## Reactivity

  Marimo's reactivity means:

  - When a variable changes, all cells that use that variable automatically re-execute
  - UI elements trigger updates when their values change without explicit callbacks
  - UI element values are accessed through \`.value\` attribute
  - You cannot access a UI element's value in the same cell where it's defined
  - Cells prefixed with an underscore (e.g. _my_var) are local to the cell and cannot be accessed by other cells

  ## Best Practices

  <data_handling>
  - Use polars for data manipulation
  - Implement proper data validation
  - Handle missing values appropriately
  - Use efficient data structures
  - A variable in the last expression of a cell is automatically displayed as a table
  </data_handling>

  <visualization>
  - For matplotlib: use plt.gca() as the last expression instead of plt.show()
  - For plotly: return the figure object directly
  - For altair: return the chart object directly. Add tooltips where appropriate. You can pass polars dataframes directly to altair.
  - Include proper labels, titles, and color schemes
  - Make visualizations interactive where appropriate
  </visualization>

  <ui_elements>
  - Access UI element values with .value attribute (e.g., slider.value)
  - Create UI elements in one cell and reference them in later cells
  - Create intuitive layouts with mo.hstack(), mo.vstack(), and mo.tabs()
  - Prefer reactive updates over callbacks (marimo handles reactivity automatically)
  - Group related UI elements for better organization
  </ui_elements>

  <sql>
  - When writing duckdb, prefer using marimo's SQL cells, which start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
  - See the SQL with duckdb example for an example on how to do this
  - Don't add comments in cells that use mo.sql()
  - Consider using \`vega_datasets\` for common example datasets
  </sql>

  ## Troubleshooting

  Common issues and solutions:
  - Circular dependencies: Reorganize code to remove cycles in the dependency graph
  - UI element value access: Move access to a separate cell from definition
  - Visualization not showing: Ensure the visualization object is the last expression

  ## Available UI elements

  - \`mo.ui.altair_chart(altair_chart)\`
  - \`mo.ui.button(value=None, kind='primary')\`
  - \`mo.ui.run_button(label=None, tooltip=None, kind='primary')\`
  - \`mo.ui.checkbox(label='', value=False)\`
  - \`mo.ui.date(value=None, label=None, full_width=False)\`
  - \`mo.ui.dropdown(options, value=None, label=None, full_width=False)\`
  - \`mo.ui.file(label='', multiple=False, full_width=False)\`
  - \`mo.ui.number(value=None, label=None, full_width=False)\`
  - \`mo.ui.radio(options, value=None, label=None, full_width=False)\`
  - \`mo.ui.refresh(options: List[str], default_interval: str)\`
  - \`mo.ui.slider(start, stop, value=None, label=None, full_width=False, step=None)\`
  - \`mo.ui.range_slider(start, stop, value=None, label=None, full_width=False, step=None)\`
  - \`mo.ui.table(data, columns=None, on_select=None, sortable=True, filterable=True)\`
  - \`mo.ui.text(value='', label=None, full_width=False)\`
  - \`mo.ui.text_area(value='', label=None, full_width=False)\`
  - \`mo.ui.data_explorer(df)\`
  - \`mo.ui.dataframe(df)\`
  - \`mo.ui.plotly(plotly_figure)\`
  - \`mo.ui.tabs(elements: dict[str, mo.ui.Element])\`
  - \`mo.ui.array(elements: list[mo.ui.Element])\`
  - \`mo.ui.form(element: mo.ui.Element, label='', bordered=True)\`

  ## Layout and utility functions

  - \`mo.md(text)\` - display markdown
  - \`mo.stop(predicate, output=None)\` - stop execution conditionally
  - \`mo.output.append(value)\` - append to the output when it is not the last expression
  - \`mo.output.replace(value)\` - replace the output when it is not the last expression
  - \`mo.Html(html)\` - display HTML
  - \`mo.image(image)\` - display an image
  - \`mo.hstack(elements)\` - stack elements horizontally
  - \`mo.vstack(elements)\` - stack elements vertically
  - \`mo.tabs(elements)\` - create a tabbed interface



  ## Examples

  <example title="Markdown ccell">
  ${formatCells([
    `
  mo.md("""
  # Hello world
  This is a _markdown_ **cell**.
  """)
    `,
  ])}
  </example>

  <example title="Basic UI with reactivity">
  ${formatCells([
    `
  import marimo as mo
  import altair as alt
  import polars as pl
  import numpy as np
    `,
    `
  n_points = mo.ui.slider(10, 100, value=50, label="Number of points")
  n_points
    `,
    `
  x = np.random.rand(n_points.value)
  y = np.random.rand(n_points.value)

  df = pl.DataFrame({"x": x, "y": y})

  chart = alt.Chart(df).mark_circle(opacity=0.7).encode(
      x=alt.X('x', title='X axis'),
      y=alt.Y('y', title='Y axis')
  ).properties(
      title=f"Scatter plot with {n_points.value} points",
      width=400,
      height=300
  )

  chart
    `,
  ])}
  </example>

  <example title="Data explorer">
  ${formatCells([
    `
  import marimo as mo
  import polars as pl
  from vega_datasets import data
    `,
    `
  cars_df = pl.DataFrame(data.cars())
  mo.ui.data_explorer(cars_df)
    `,
  ])}
  </example>

  <example title="Multiple UI elements">
  ${formatCells([
    `
  import marimo as mo
  import polars as pl
  import altair as alt
    `,
    `
  iris = pl.read_csv("hf://datasets/scikit-learn/iris/Iris.csv")
    `,
    `
  species_selector = mo.ui.dropdown(
      options=["All"] + iris["Species"].unique().to_list(),
      value="All",
      label="Species",
  )
  x_feature = mo.ui.dropdown(
      options=iris.select(pl.col(pl.Float64, pl.Int64)).columns,
      value="SepalLengthCm",
      label="X Feature",
  )
  y_feature = mo.ui.dropdown(
      options=iris.select(pl.col(pl.Float64, pl.Int64)).columns,
      value="SepalWidthCm",
      label="Y Feature",
  )
  mo.hstack([species_selector, x_feature, y_feature])
    `,
    `
  filtered_data = iris if species_selector.value == "All" else iris.filter(pl.col("Species") == species_selector.value)

  chart = alt.Chart(filtered_data).mark_circle().encode(
      x=alt.X(x_feature.value, title=x_feature.value),
      y=alt.Y(y_feature.value, title=y_feature.value),
      color='Species'
  ).properties(
      title=f"{y_feature.value} vs {x_feature.value}",
      width=500,
      height=400
  )

  chart
    `,
  ])}
  </example>

  <example title="Conditional Outputs">
  ${formatCells([
    `
  mo.stop(not data.value, mo.md("No data to display"))

  if mode.value == "scatter":
    mo.output.replace(render_scatter(data.value))
  else:
    mo.output.replace(render_bar_chart(data.value))
    `,
  ])}
  </example>

  <example title="Interactive chart with Altair">
  ${formatCells([
    `
  import marimo as mo
  import altair as alt
  import polars as pl
    `,
    `# Load dataset
  weather = pl.read_csv("https://raw.githubusercontent.com/vega/vega-datasets/refs/heads/main/data/weather.csv")
  weather_dates = weather.with_columns(
      pl.col("date").str.strptime(pl.Date, format="%Y-%m-%d")
  )
  _chart = (
      alt.Chart(weather_dates)
      .mark_point()
      .encode(
          x="date:T",
          y="temp_max",
          color="location",
      )
  )
  `,
    "chart = mo.ui.altair_chart(_chart)\nchart",
    `
  # Display the selection
  chart.value`,
  ])}
  </example>

  <example title="Run Button Example">
  ${formatCells([
    "import marimo as mo",
    `
  first_button = mo.ui.run_button(label="Option 1")
  second_button = mo.ui.run_button(label="Option 2")
  [first_button, second_button]`,
    `
  if first_button.value:
      print("You chose option 1!")
  elif second_button.value:
      print("You chose option 2!")
  else:
      print("Click a button!")`,
  ])}
  </example>

  <example title="SQL with duckdb">
  ${formatCells([
    "import marimo as mo\n  import polars as pl",
    `weather = pl.read_csv('https://raw.githubusercontent.com/vega/vega-datasets/refs/heads/main/data/weather.csv')`,
    `seattle_weather_df = mo.sql(
      f"""
      SELECT * FROM weather WHERE location = 'Seattle';
      """
  )`,
  ])}
  </example>`;
}

function formatCells(cells: string[]) {
  const indent = "    ";
  const indentCode = (code: string) => {
    return code
      .trim()
      .split("\n")
      .map((line) => (indent + line).trimEnd())
      .join("\n");
  };

  const formatCell = (cell: string) => {
    return `
  @app.cell
  def _():
  ${indentCode(cell)}
      return
  `;
  };

  return cells.map(formatCell).join("");
}