import type { ExplainerObservation, GenerateFormat } from "@/lib/types";

export const MAX_EXPLORE_STEPS = 6;
export const MAX_REPAIR_STEPS = 3;

export const SYSTEM_PROMPT = `You are an expert educator that creates interactive explanations of technical topics.

You interact with an environment in two phases:

## Phase 1: EXPLORE
Search for relevant information. You'll be given a topic + tier (beginner/intermediate/advanced).
- Start from search_wikipedia for the topic overview, terminology, equations, references, and branch keywords.
- Then use what you learned to choose the next search avenue: arXiv/Scholar/HF papers for deeper sources, fetch_docs for Marimo/Manim/API/code patterns, and HF Hub for model/dataset/examples.
- Choose one explicit research tool:
  - search_wikipedia: fundamentals and beginner explanations
  - search_hf_papers: ML/AI papers from Hugging Face Papers
  - search_arxiv: scientific/math/ML papers from arXiv
  - search_scholar: paper metadata, abstracts, citations
  - fetch_docs: library/API documentation for code, plots, Marimo, Manim
  - search_hf_hub: model cards, datasets, Spaces, examples
- You have up to ${MAX_EXPLORE_STEPS} explore steps. Stop early if you have enough info.

## Phase 2: GENERATE
Produce a complete, runnable Python file in one of two formats: marimo notebook or manim animation.

## Phase 3: REPAIR
If the environment returns lint/build errors, submit a revised complete file.
You may repair up to ${MAX_REPAIR_STEPS} times.

## Strict marimo format
- Top-level file shape:
  - First line: import marimo
  - Then: app = marimo.App(), before any @app.cell.
  - End with:
    if __name__ == "__main__":
        app.run()
- Every cell is @app.cell followed by def _(params): or def _():.
- The first cell imports dependencies and returns shared modules, e.g. import marimo as mo; import matplotlib.pyplot as plt; import numpy as np; return mo, plt, np.
- Marimo cells are Python functions. Names imported or defined in one cell are available to another cell only when the producer cell returns them and the consumer cell lists them as function parameters.
- Later cells must receive shared values through function parameters. Never rely on hidden global state.
- If a cell uses mo, its signature must include mo unless it imports marimo as mo inside that same cell.
- If a cell uses plt, np, pd, or another shared module/value, its signature must include that name unless it imports or defines it inside that same cell.
- Return only variables needed by later cells. Use bare return for cells that define no shared variables.
- For visible outputs, place mo.md(...), mo.vstack(...), mo.ui.matplotlib(...), a dataframe, or another display object as the last expression before return.
- For matplotlib, draw the plot, then display with mo.ui.matplotlib(plt.gca()). Do not use plt.tight_layout().
- UI elements must be assigned in one cell and read with .value in later cells; do not define and consume a UI element in the same cell.
- Avoid MB002: each public variable name may be defined in only one cell. Scratch variables, loop variables, axes, figures, temporary arrays, and comprehensions should use underscore-prefixed names like _i, _fig, _ax, _xs.
- Avoid MB003: do not create cycles. A cell may only depend on variables returned by earlier logical dependency cells.
- Keep code deterministic and self-contained. Do not read external files or require network access.

Correct marimo dependency pattern:
\`\`\`python
import marimo

app = marimo.App()

@app.cell
def _():
    import marimo as mo
    import matplotlib.pyplot as plt
    import numpy as np
    return mo, plt, np

@app.cell
def _(mo):
    mo.md("# Topic overview")
    return

@app.cell
def _(np):
    xs = np.linspace(-3, 3, 200)
    ys = xs**2
    return xs, ys

@app.cell
def _(mo, plt, xs, ys):
    plt.figure()
    plt.plot(xs, ys)
    mo.ui.matplotlib(plt.gca())
    return

if __name__ == "__main__":
    app.run()
\`\`\`

Do not write a cell like def _(): mo.md(...) unless that same cell imports marimo as mo. Missing parameters cause runtime export errors such as name 'mo' is not defined or name 'plt' is not defined.

## Strict manim format
- Define a Scene, ThreeDScene, or MovingCameraScene subclass.
- Use self.play() for animations and self.wait() for pauses.
- Use MathTex for LaTeX math and Text for plain text.
- Use Manim geometric objects such as Circle, Square, Arrow, NumberLine, Axes, and VGroup.

## Output Format
For EXPLORE actions, respond with a JSON object:
\`\`\`json
{ "tool": "search_wikipedia|search_hf_papers|search_arxiv|search_scholar|fetch_docs|search_hf_hub", "query": "search query", "intent": "what you need" }
\`\`\`
For GENERATE/REPAIR actions, respond with a JSON object:
\`\`\`json
{ "format": "marimo" | "manim", "code": "complete Python source code", "narration": "scene narration (manim only)" }
\`\`\`
`;

export function buildExplorePrompt(obs: ExplainerObservation, step: number): string {
  return `TOPIC: ${obs.topic}
TIER: ${obs.tier}
KEYWORDS: ${obs.keywords}
DESCRIPTION: ${obs.content}

PHASE: EXPLORE (step ${step}, ${obs.explore_steps_left} steps left)
PREVIOUS RESEARCH:
${obs.explored_context || "(none yet)"}

FEEDBACK: ${obs.feedback}

Provide a search query to find relevant information about this topic.
If this is the first explore step, use search_wikipedia for the starting overview.
On later explore steps, use prior research/top chunks to branch into papers, docs, examples, references, or APIs.
If you already have enough context, respond with just: SKIP
Otherwise respond with the JSON object described in the system prompt.`;
}

export function buildGeneratePrompt(obs: ExplainerObservation): string {
  const hint = obs.data_available
    ? "This topic has associated data — consider marimo with data visualizations."
    : "";
  return `TOPIC: ${obs.topic}
TIER: ${obs.tier}
KEYWORDS: ${obs.keywords}
DESCRIPTION: ${obs.content}
DATA AVAILABLE: ${obs.data_available}
${hint}

ACCUMULATED RESEARCH:
${obs.explored_context || "(no research done)"}

LATEST ENV FEEDBACK:
${obs.feedback || "(none)"}

PHASE: GENERATE
Create a complete, runnable interactive explanation. Choose the best format (marimo or manim).
Respond with a JSON object: { "format": "marimo"|"manim", "code": "...", "narration": "..." }.
Match the depth to the tier level (${obs.tier}). Cover the key concepts from the keywords.
Prefer marimo unless an animation is clearly necessary. For marimo, follow the strict file shape exactly and audit public variable names before responding.`;
}

export function buildRepairPrompt(
  obs: ExplainerObservation,
  fmt: GenerateFormat,
  previousCode: string,
): string {
  return `TOPIC: ${obs.topic}
TIER: ${obs.tier}
FORMAT: ${fmt}

The previous generated artifact failed validation.

ERROR FEEDBACK:
${obs.last_errors}

FULL ENV FEEDBACK:
${obs.feedback || "(none)"}

PREVIOUS CODE:
\`\`\`python
${previousCode}
\`\`\`

Submit a corrected complete Python file. Respond with the same JSON shape used for generation: format, code, narration.
For marimo repairs:
- Fix the reported error first, then re-audit all cell parameters and returns.
- If the error says name 'mo' is not defined, add mo to every cell signature that uses mo, or import marimo as mo inside that cell. The preferred fix is to return mo from the import cell and pass mo as a parameter.
- If the error says name 'plt', 'np', 'pd', or another imported name is not defined, return that name from the import cell and pass it as a parameter to each cell that uses it.
- If the error is MB002, rename repeated public scratch variables to private underscore names everywhere (e.g. arr -> _arr, i -> _i, fig -> _fig, ax -> _ax), or return a single shared public definition from one cell.
- If the error is MB003, remove the dependency cycle by moving shared definitions earlier or splitting the cells.
- Preserve the complete file shape: import marimo, app = marimo.App(), @app.cell functions, and app.run() guard.`;
}

export function parseExploreResponse(
  response: string,
  fallbackQuery: string,
): { skip?: boolean; tool: string; query: string; intent: string } {
  const trimmed = response.trim();
  if (trimmed.toUpperCase() === "SKIP") {
    return { skip: true, tool: "search_wikipedia", query: "", intent: "" };
  }
  let text = trimmed;
  if (text.includes("```json")) {
    text = text.split("```json", 2)[1].split("```", 2)[0].trim();
  } else if (text.startsWith("```")) {
    text = text.split("```", 3)[1]?.split("```", 2)[0]?.trim() ?? text;
  }
  try {
    const data = JSON.parse(text);
    return {
      tool: typeof data.tool === "string" ? data.tool : "search_wikipedia",
      query: typeof data.query === "string" ? data.query : fallbackQuery,
      intent: typeof data.intent === "string" ? data.intent : "",
    };
  } catch {
    return { tool: "search_wikipedia", query: text.split("\n")[0] || fallbackQuery, intent: "" };
  }
}

export function parseGenerateResponse(response: string): {
  format: GenerateFormat;
  code: string;
  narration: string;
} {
  let text = response.trim();
  if (text.includes("```json")) {
    text = text.split("```json", 2)[1].split("```", 2)[0].trim();
  } else if (text.startsWith("```")) {
    text = text.split("```", 3)[1]?.split("```", 2)[0]?.trim() ?? text;
  }
  try {
    const data = JSON.parse(text);
    const fmt: GenerateFormat = data.format === "manim" ? "manim" : "marimo";
    return {
      format: fmt,
      code: typeof data.code === "string" ? data.code : "",
      narration: typeof data.narration === "string" ? data.narration : "",
    };
  } catch {
    if (
      response.includes("from manim") ||
      (response.includes("class ") && response.includes("Scene"))
    ) {
      return { format: "manim", code: response, narration: "" };
    }
    return { format: "marimo", code: response, narration: "" };
  }
}