"""LLM prompt templates for BTQL query refinement."""

BTQL_REFERENCE = """
## BTQL Syntax Reference

BTQL uses pipe-delimited clauses. The parser auto-detects BTQL vs SQL based on whether query starts with `select:`, `from:`, `filter:`, etc.

### Basic Structure
```
from: experiment('<experiment_id>')
select: field1, field2
filter: <conditions>
dimensions: field1, field2  -- equivalent to GROUP BY
measures: count(1), avg(scores.metric)  -- aggregate functions
sort: field desc
limit: 100
```

### Data Sources
- `experiment('<id>')` - query a specific experiment
- `project_logs('<id>')` - query project logs
- Shape options: `spans` (default), `traces`, `summary`
  - Example: `from: experiment('id') summary`

### Field Access
- Object fields: `metadata.model`, `scores.Factuality`
- Fields with spaces/special chars: `scores."clarity & readability"`
- Array access: `tags[0]`, `tags[-1]` (last element)
- Nested: `input.output.picasso_eval_metadata.inferences[-1].response`

### Operators
- Comparison: `=`, `!=`, `>`, `<`, `>=`, `<=`
- Null checks: `IS NULL`, `IS NOT NULL`
- Logical: `AND`, `OR`, `NOT` (case-insensitive)
- Text: `LIKE`, `NOT LIKE`, `ILIKE` (case-insensitive), `MATCH`
- Array: `INCLUDES`, `NOT INCLUDES`, `contains`, `not contains`
- Arithmetic: `+`, `-`, `*`, `/`, `%`

### Functions
- Date: `day(created)`, `hour(created)`, `month(created)`, `now()`
- String: `lower()`, `upper()`, `concat()`
- Array: `len()`, `contains()`
- Null handling: `coalesce()`, `nullif()`
- Aggregates (use in measures): `count()`, `sum()`, `avg()`, `min()`, `max()`, `percentile(field, p)`

### Common Patterns

**Fetch scores from experiment:**
```
from: experiment('<id>')
select: root_span_id, scores
filter: scores IS NOT NULL
```

**Aggregate scores:**
```
from: experiment('<id>') summary
dimensions: root_span_id
measures: avg(scores.completeness), avg(scores."clarity & readability")
filter: scores IS NOT NULL
```

**With custom filters:**
```
from: experiment('<id>')
select: root_span_id, scores
filter: scores IS NOT NULL and metadata.model = "gpt-4"
```
"""

SYSTEM_PROMPT = f"""You are a BTQL (Braintrust Query Language) expert. Your task is to write or fix BTQL queries to fetch metrics from Braintrust experiments.

{BTQL_REFERENCE}

## Important Notes
1. Metric names with spaces or special characters MUST be quoted: `scores."clarity & readability"`
2. Use `summary` shape when aggregating across spans: `from: experiment('id') summary`
3. For fetching raw scores, use `select:` with `scores` field
4. For aggregated metrics, use `dimensions:` and `measures:`
5. Always include `filter: scores IS NOT NULL` to exclude empty results
6. The `AND`/`OR`/`NOT` operators are case-insensitive

When you receive an error, analyze it carefully and fix the query. Common issues:
- Unquoted field names with special characters
- Wrong field paths
- Syntax errors in filter expressions
- Missing required clauses
"""

INITIAL_QUERY_PROMPT = """Generate a BTQL query to fetch the following metrics from an experiment.

Experiment ID: {experiment_id}
Metrics to fetch: {metrics}
Additional filters: {filters}

Requirements:
1. Query should return individual score values that can be aggregated later
2. Include root_span_id to identify unique records
3. Apply the additional filters if provided
4. Use proper quoting for metric names with special characters

Return ONLY the BTQL query, no explanation or markdown formatting.
"""

REFINEMENT_PROMPT = """The previous BTQL query failed. Analyze the error and provide a corrected query.

Experiment ID: {experiment_id}
Metrics to fetch: {metrics}
Additional filters: {filters}

Previous query:
{previous_query}

Error message:
{error_message}

Previous attempts:
{attempt_history}

Analyze the error and fix the query. Common fixes:
- Quote field names with spaces: scores."field name"
- Check field path correctness
- Verify filter syntax
- Ensure proper operator usage

Return ONLY the corrected BTQL query, no explanation or markdown formatting.
"""

SCHEMA_ANALYSIS_PROMPT = """Analyze this sample data from a Braintrust experiment and identify where metrics and explanations are stored, and also find where AI model's response/output is stored.

Sample data (1-2 rows):
{sample_data}

{attempt_history_section}

{metrics_section}

Your task:
{task_step_1}
- Search DEEPLY through ALL nested objects for explanations/feedback for every metric
- Search DEEPLY through ALL nested objects for the AI model's response/output

IMPORTANT: 
When looking for metrics, look under the "scores" field for all rows to get distinct universal set.

When searching for explanation, try to find the explanation for every metric. And the possible locations in ALL of these:
- metadata.<metric_name>.explanation (direct string)
- metadata.<metric_name>.feedback (direct string)
- input.output[] array - look for judge_response.text containing JSON with feedback/explanation fields
- output.choices[].message.content containing JSON with feedback/explanation fields
- Any field containing "explanation", "feedback", "rationale", or "reason"
- Prefer stable parent paths over brittle numeric-key leaf paths. Example: if values are under metadata.ratings."1".rationale, prefer metadata.ratings as a json source path.

When searching for the AI model's response (the actual answer/output the model generated), look CAREFULLY in these locations:
- output (direct field)
- input.output (may be a string or array)
- input.output[0] or input.output[-1] (first or last element)
- input.output.chat_context.turns[].message (iterate through turns to find the assistant's message)
- input.output.chat_context.state.turns[].message (with state wrapper)
- Any field named: "response", "ai_response", "picasso_ai_response", "test_model_response", "assistant_message", "model_output"
- Look for the LONGEST text field that appears to be the model's actual answer (not system prompts, not user inputs)
- Check both input and output top-level fields thoroughly

CRITICAL for response_field_path:
- MUST use bracket notation for array indices: "input.output[0].text" NOT "input.output.0.text"
- If the response is in an array, specify the index (usually 0 for first, or -1 for last)
- Verify the path points to actual response text, not metadata or timestamps
- If you find multiple potential paths, choose the one with the most substantive text content

For some metrics, their explanations are often stored as JSON strings that need parsing. Look for JSON objects containing keys matching the metric names with nested "feedback" or "explanation" fields.

Return a JSON object with this exact structure:
{{
    "scores_path": "scores",
    "available_metrics": ["metric1", "metric2"],
    "explanation_paths": [
        {{"metric": "metric1", "path": "metadata.metric1.explanation", "type": "direct"}},
        {{"metric": "metric2", "path": "input.output[0].judge_response.text", "type": "json", "json_key": "metric2.feedback"}}
    ],
    "response_field_path": "input.output[0].picasso_ai_response",
    "id_field": "root_span_id"
}}

Notes on explanation_paths:
- "type": "direct" means the path points directly to a string explanation
- "type": "json" means the path points to a JSON string that needs parsing, with json_key specifying the nested key to extract (e.g., "completeness.feedback" or "groundedness.explanation")
- "type": "not_found" means the explanation path for this metric was not found in the sample data
- Use array indices (0, 1, etc.) for array access in paths
- if there's a attribute didn't follow the program variable naming rules, please use double quotes to wrap it
- Include ALL metrics found in scores, even if you can't find their explanation paths
- Do NOT copy example paths blindly; every path must be grounded in the provided sample data.
- Avoid fragile object numeric-key leaf paths like metadata.ratings.1.rationale; use a stable parent json path when needed.

Notes on response_field_path:
- MUST use dot notation with numeric indices: "input.output[0].text" NOT "input.output.0.text"
- Should point to the actual text content of the AI's response
- Common patterns: "output", "input.output[0]", "input.output.chat_context.turns[-1]message"
- Use -1 for last element in arrays when appropriate
- This is the most important field - take extra time to find the correct path by examining the sample data

Return ONLY the JSON object, no explanation or markdown formatting.
"""

EXPLANATION_ANALYSIS_PROMPT = """Analyze this sample data from a Braintrust experiment and identify where metric explanations are stored.

Sample data (1-2 rows):
{sample_data}

{attempt_history_section}

Metrics to find explanations for:
{metrics}

Your task:
- Search DEEPLY through ALL nested objects for explanations/feedback for every metric
- Find the MOST SPECIFIC path that points directly to the explanation content
- Check ALL of these possible locations:
  - metadata.<metric_name>.explanation (direct string)
  - metadata.<metric_name>.feedback (direct string)
  - metadata.<metric_name>.rationale (direct string)
  - metadata.<metric_name>.turn_0_judge_response (or similar turn_*_judge_response patterns)
  - metadata.<metric_name>.<metric_name>_judge_tips.explanation
  - metadata.reasoning, metadata.rationale (generic explanation fields)
  - metadata.ratings (may contain numbered objects like 1, 2, 3 with rationale fields inside)
  - Any field containing "explanation", "feedback", "rationale", "reason", "ratings", "judge_response" etc.
  - do not use [*] or .* notation for arrays or objects.
  - sometimes explanations value may be null or empty, but the key indicates that it's very likely the path of the explanation, you still need to record the path.
  - prefer stable parent paths over brittle numeric-key leaf paths (e.g., metadata.ratings instead of metadata.ratings.1.rationale)

CRITICAL - Use the MOST FINE-GRAINED path possible:
- DO NOT stop at container objects if a more specific path exists
- If metadata.groundedness contains both "turn_0_judge_response" (explanation) and "turn_0_picasso_response_text" (model response), use "metadata.groundedness.turn_0_judge_response" NOT "metadata.groundedness"
- Always drill down to the exact field containing the explanation/feedback
- Avoid paths that would include sibling fields that are NOT explanations
- Example: metadata.groundedness.turn_0_judge_response is BETTER than metadata.groundedness
- Example: metadata.completeness.feedback is BETTER than metadata.completeness

IMPORTANT - Don't assume field names match metric names:
- "percentage_of_instructions_followed" might have explanations in metadata.ratings (a container with numbered sub-objects)
- "groundedness" might have explanations in metadata.groundedness.turn_0_judge_response or metadata.groundedness.groundedness_judge_tips.explanation
- "completeness" might have explanations in metadata.reasoning or metadata.completeness.feedback
- Search the ENTIRE metadata structure, not just fields matching metric names
- When you find a container object with numbered keys (like ratings: {{1: {{rationale: ...}}, 2: {{rationale: ...}}}}), record the parent path (e.g., "metadata.ratings")

For complex nested structures:
- If explanations are in numbered objects (like metadata.ratings.1.rationale, metadata.ratings.2.rationale), just record "metadata.ratings" as the path
- The json_key can be used to indicate which sub-field contains the actual explanation text (e.g., "rationale")

For some metrics, their explanations are often stored as JSON strings that need parsing. Look for JSON objects containing keys matching the metric names with nested "feedback" or "explanation" fields.

Return a JSON object with this exact structure:
{{
    "scores_path": "scores",
    "available_metrics": ["metric1", "metric2"],
    "explanation_paths": [
        {{"metric": "metric1", "path": "metadata.metric1.explanation", "type": "direct"}},
        {{"metric": "metric2", "path": "input.output[0].judge_response.text", "type": "json", "json_key": "metric2.feedback"}},
        {{"metric": "groundedness", "path": "metadata.groundedness.turn_0_judge_response", "type": "json", "json_key": "groundedness.feedback"}},
        {{"metric": "percentage_of_instructions_followed", "path": "metadata.ratings", "type": "direct", "json_key": "rationale"}}
    ],
    "id_field": "root_span_id"
}}

Notes on explanation_paths:
- "type": "direct" means the path points directly to a string explanation (or to a structure containing explanations)
- "type": "json" means the path points to a JSON string that needs parsing, with json_key specifying the nested key to extract
- "type": "not_found" means the explanation path for this metric was not found in the sample data
- "json_key" can specify which field to extract from complex nested objects (e.g., "rationale" when path is "metadata.ratings")
- Use bracket notation for array indices: "input.output[0].text" NOT "input.output.0.text"
- For field names, use DOT NOTATION by default: "metadata.factuality_rating.rationale" NOT "metadata["factuality_rating"].rationale"
- ONLY use double quotes for field names that don't follow standard variable naming rules (contain spaces, colons, dashes, or start with numbers):
  - CORRECT: metadata."Bing:#-searchable-claims".raw_claims (special characters)
  - CORRECT: metadata."clarity & readability".feedback (has spaces)
  - CORRECT: metadata.factuality_rating.rationale (standard name, use dots)
  - WRONG: metadata["factuality_rating"].rationale (don't use brackets for standard names)
- Include ALL metrics, even if you can't find their explanation paths (mark as "not_found")
- Do NOT copy example paths blindly; every path must be grounded in the provided sample data.
- Avoid fragile object numeric-key leaf paths like metadata.ratings.1.rationale; use a stable parent json path when needed.

Return ONLY the JSON object, no explanation or markdown formatting.
"""

RESPONSE_ANALYSIS_PROMPT = """Analyze this sample data from a Braintrust experiment and identify where the AI model's response/output is stored.

Sample data (1-2 rows):
{sample_data}

{attempt_history_section}

Your task:
- Search DEEPLY through ALL nested objects for the AI model's response/output
- Look CAREFULLY in these locations:
  - output (direct field)
  - input.output (may be a string or array)
  - input.output[0] or input.output[-1] (first or last element)
  - input.output.chat_context.turns[].message (iterate through turns to find the assistant's message)
  - input.output.chat_context.state.turns[].message (with state wrapper)
  - Any field named: "response", "ai_response", "picasso_ai_response", "test_model_response", "assistant_message", "model_output"
  - Look for the LONGEST text field that appears to be the model's actual answer (not system prompts, not user inputs)
  - Check both input and output top-level fields thoroughly

CRITICAL for response_field_path:
- MUST use bracket notation for array indices: "input.output[0].text" NOT "input.output.0.text"
- If the response is in an array, specify the index (usually 0 for first, or -1 for last)
- Verify the path points to actual response text, not metadata or timestamps
- If you find multiple potential paths, choose the one with the most substantive text content
- If NO valid path is found after thorough search, return null for response_field_path

Return a JSON object with this exact structure:
{{
    "response_field_path": "input.output[0].picasso_ai_response"
}}

Or if no path found:
{{
    "response_field_path": null
}}

Notes on response_field_path:
- Should point to the actual text content of the AI's response
- Common patterns: "output", "input.output[0]", "input.output.chat_context.turns[-1].message"
- Use -1 for last element in arrays when appropriate
- Return null if you cannot find a valid path after exhaustive search
- This is the most important field - take extra time to find the correct path by examining the sample data

Return ONLY the JSON object, no explanation or markdown formatting.
"""

EXPLANATION_VERIFICATION_PROMPT = """You are verifying query results from a Braintrust experiment for explanation extraction.

The query was designed to fetch explanations/feedback for LLM-judged metrics.

Here is a sample result row from executing the query:
{sample_result}

Here are the target metrics for explanation extraction:
{available_metrics}

Verify the following:

1. **judge_response field**: Does it contain meaningful explanation/feedback text?
   - judge_response may be plain text or JSON-like text
   - It is valid if at least ONE target metric has useful explanation evidence
   - Do NOT require every listed metric to appear in this single row
   - Empty, null, "{{}}", or unrelated operational text is NOT valid

2. **input field**: Does it contain meaningful input?
   - Should contain the user's question or request

Return a JSON object with this structure:
{{
    "is_valid": true/false,
    "explanation": "Brief explanation of why the row is valid or what's wrong with judge_response"
}}

Return ONLY the JSON object, no markdown formatting.
"""

RESPONSE_VERIFICATION_PROMPT = """You are verifying query results from a Braintrust experiment for AI model response extraction.

The query was designed to fetch the AI model's actual response/output.

Here is a sample result row from executing the query:
{sample_result}

Verify the following:

1. **test_model_response field**: Does it contain the AI model's actual response?
   - Should be substantive text content that looks like an AI-generated answer
   - Should NOT be null, empty, metadata, or system info
   - Should look like a response to the `input` field
   - Should have meaningful content, not just timestamps or status codes

2. **input field**: Does it contain meaningful input?
   - Should contain the user's question or request

Return a JSON object with this structure:
{{
    "is_valid": true/false,
    "explanation": "Brief explanation of why the data is valid or what's wrong with test_model_response"
}}

Return ONLY the JSON object, no markdown formatting.
"""

QUERY_TEMPLATE_VERIFICATION_PROMPT = """You are verifying query results from a Braintrust experiment.

The query was designed to fetch:
1. `judge_response` - Should contain explanations/feedback for metric scores
2. `test_model_response` - Should contain the AI model's actual response/output  
3. `input` - Should contain the user's input query given to the model

Here is a sample result row from executing the query:
{sample_result}

Here are the metrics this query should support:
{available_metrics}

Verify the following:

1. **judge_response field**: Does it contain explanations/feedback for the metrics?
   - Should contain score explanations or feedback text (may be JSON format)
   - Check if it has content related to the metrics listed above
   - Empty, null, or unrelated content is NOT valid

2. **test_model_response field**: Does it contain the AI model's actual response?
   - Should be substantive text content that looks like an AI-generated answer
   - Should NOT be null, empty, metadata, or system info
   - Should look like a response to the `input` field

3. **input field**: Does it contain meaningful input?
   - Should contain the user's question or request

Return a JSON object with this structure:
{{
    "is_valid": true/false,
    "explanation": "Brief explanation of why the data is valid or what's wrong with each field"
}}

Return ONLY the JSON object, no markdown formatting.
"""

EXPLANATION_PATH_SELECTION_PROMPT = """You are selecting the best explanation path configuration from multiple candidates.

All candidates successfully executed without query errors, but none passed full verification.
Your task is to select the BEST candidate based on:
1. Coverage: How many of the required metrics have explanations extracted
2. Quality: How meaningful the extracted explanations appear to be

Required metrics to cover: {available_metrics}

Here are the candidates (each shows the explanation paths discovered and a sample of extracted data):

{candidates_info}

Evaluate each candidate:
1. Count how many required metrics have non-null, non-empty explanations in the sample
2. Check if the explanations look like actual feedback/reasoning (not just null, empty strings, or metadata)
3. Prioritize candidates with higher metric coverage

Return a JSON object with this structure:
{{
    "selected_index": <0-based index of the best candidate>,
    "explanation": "Brief explanation of why this candidate was selected",
    "metrics_coverage": {{
        "selected_candidate_metrics_found": ["list", "of", "metrics", "found"],
        "coverage_percentage": <percentage as integer>
    }}
}}

Return ONLY the JSON object, no markdown formatting.
"""
