import urllib.request
import pathlib
import pandas as pd
import duckdb

# --- Download the raw data files ---
DATA_DIR = pathlib.Path("data")
DATA_DIR.mkdir(exist_ok=True)
BASE_URL = "https://raw.githubusercontent.com/jpatokal/openflights/master/data/"
FILES_TO_DOWNLOAD = {
    "airports": "airports.dat",
    "airlines": "airlines.dat",
    "routes": "routes.dat",
    "countries": "countries.dat",
    "planes": "planes.dat"
}
# Define column names as the files don't have headers
COLUMN_NAMES = {
    "airports": ["airport_id", "name", "city", "country", "iata", "icao", "latitude", "longitude", "altitude", "timezone", "dst", "tz_db", "type", "source"],
    "airlines": ["airline_id", "name", "alias", "iata", "icao", "callsign", "country", "active"],
    "routes": ["airline", "airline_id", "source_airport", "source_airport_id", "destination_airport", "destination_airport_id", "codeshare", "stops", "equipment"],
    "countries": ["name", "iso_code", "dafif_code"],
    "planes": ["name", "iata", "icao"]
}

PROD_DB_PATH = "data/prod_openflights.db"

# --- Load the real data into our "production" DuckDB ---
with duckdb.connect(PROD_DB_PATH) as con:
    for name, filename in FILES_TO_DOWNLOAD.items():
        url = f"{BASE_URL}{filename}"
        path = DATA_DIR / filename
        if not path.exists():
            urllib.request.urlretrieve(url, path)
            print(f"✅ Downloaded: {path}")

        # Load data using pandas to handle missing headers and null values
        df = pd.read_csv(path, header=None, names=COLUMN_NAMES[name], na_values=["\\N"])
        con.execute(f"CREATE OR REPLACE TABLE {name} AS SELECT * FROM df")

    print(f"\n✅ 'Production' database simulated at: {PROD_DB_PATH}")
    print("Tables created:", con.sql("SHOW TABLES;").fetchall())

import duckdb

# Connect to the "production" database we just created
with duckdb.connect(PROD_DB_PATH, read_only=True) as con:
    # The DESCRIBE command gives us the schema information for all tables
    schema_df = con.sql("DESCRIBE;").df()

print("✅ Schema successfully extracted from 'production' database:")
print(schema_df.to_markdown(index=False))

# We can also store this for later use in prompts
schema_for_prompt = schema_df.to_markdown(index=False)

import pandas as pd
import os
from pydantic import create_model, BaseModel
from fireworks import LLM
import duckdb
import json
from dotenv import load_dotenv
from typing import List, Optional, Any, Dict, Type
import datetime
import decimal
import uuid
import math
import time


TARGET_ROW_COUNT = 100  # The number of rows to generate for each table.

# --- 1. Dynamically Create Pydantic Models from the SQL Schema ---
def map_sql_type_to_python(sql_type: str) -> Type:
    """Maps SQL data types to Python types for Pydantic models."""
    sql_type_upper = str(sql_type).upper()
    if 'DECIMAL' in sql_type_upper: return decimal.Decimal
    if 'DOUBLE' in sql_type_upper or 'FLOAT' in sql_type_upper or 'REAL' in sql_type_upper: return float
    if 'BIGINT' in sql_type_upper or 'INT' in sql_type_upper: return int
    if 'VARCHAR' in sql_type_upper or 'TEXT' in sql_type_upper or 'STRING' in sql_type_upper: return str
    if 'TIMESTAMP' in sql_type_upper: return datetime.datetime
    if 'DATE' in sql_type_upper: return datetime.date
    if 'TIME' in sql_type_upper: return datetime.time
    if 'BOOLEAN' in sql_type_upper: return bool
    if 'BLOB' in sql_type_upper or 'BYTEA' in sql_type_upper: return bytes
    if 'UUID' in sql_type_upper: return uuid.UUID
    return object

pydantic_models: Dict[str, Type[BaseModel]] = {}
table_names = schema_df['name'].unique()

for table_name in table_names:
    table_schema = schema_df[schema_df['name'] == table_name].iloc[0]
    fields: Dict[str, Any] = {}
    col_names = table_schema['column_names']
    col_types = table_schema['column_types']
    for i, col_name in enumerate(col_names):
        python_type = map_sql_type_to_python(col_types[i])
        fields[col_name] = (Optional[python_type], None)
    model_name = table_name.capitalize() + "Model"
    pydantic_models[table_name] = create_model(model_name, **fields)

dataset_fields: Dict[str, Any] = {
    table_name: (List[model], ...) for table_name, model in pydantic_models.items()
}
SyntheticDataset = create_model('SyntheticDataset', **dataset_fields)
print("✅ Dynamically created Pydantic models for all tables.")


# --- 2. Define Total Row Counts and Chunking Strategy ---
TOTAL_ROW_COUNTS = {name: TARGET_ROW_COUNT for name in table_names}
ROWS_PER_API_CALL = 2 # Ask for data in small, safe chunks
print("\n✅ Data Generation Plan:")
print(f" - Target rows per table: {list(TOTAL_ROW_COUNTS.values())[0]}")
print(f" - Will make API calls asking for {ROWS_PER_API_CALL} rows/call until target is met.")


# --- 3. Setup LLM and Loop to Generate Data in Chunks ---
SYNTHETIC_DB_PATH = "data/synthetic_openflights.db"
load_dotenv()
llm = LLM(model="accounts/fireworks/models/deepseek-v3", deployment_type="serverless", api_key=os.getenv("FIREWORKS_API_KEY"))

all_synthetic_data: Dict[str, List[Dict]] = {name: [] for name in table_names}
chunk_row_counts = {name: ROWS_PER_API_CALL for name in table_names}

base_generation_prompt = f"""
You are a highly intelligent AI data generator. Your task is to create a realistic, synthetic dataset based on the provided database schema.
The data you generate must be internally consistent. For example, an `airline_id` in a `routes` table must correspond to an existing `airline_id` in an `airlines` table within this same generated chunk.
This applies to any schema you might be working with, not just airline-related data.
You must generate a single JSON object that strictly adheres to the provided JSON schema.

The database schema is as follows:
{schema_for_prompt}
"""

call_count = 0
# Loop until all tables have at least the desired number of rows
while not all(len(rows) >= TOTAL_ROW_COUNTS[name] for name, rows in all_synthetic_data.items()):
    call_count += 1
    print(f"\n📞 --- Generating data chunk #{call_count} ---")
    
    # --- Create a summary of existing data to guide the LLM ---
    existing_data_summary = ""
    if any(len(rows) > 0 for rows in all_synthetic_data.values()):
        summary_parts = ["\nYou have already generated the following data. Do NOT generate rows that are substantially similar to these examples. Create new, unique data.\n"]
        for table_name, rows in all_synthetic_data.items():
            if rows:
                summary_parts.append(f"\n--- Existing data in '{table_name}' table ---")
                sample_rows = rows[-100:] if len(rows) > 100 else rows  # sample the last 100 rows
                df = pd.DataFrame(sample_rows)
                if len(df.columns) > 10:
                    df = df.iloc[:, :10]
                markdown_summary = df.to_markdown(index=False, tablefmt="grid")
                if markdown_summary:
                    summary_parts.append(markdown_summary)
        existing_data_summary = "\n".join(summary_parts)


    # --- Construct the final prompt for this iteration ---
    final_prompt = (
        base_generation_prompt +
        existing_data_summary +
        f"\n\nNow, generate a NEW JSON object with a key for each table. The number of new rows for each table should be:\n" +
        json.dumps(chunk_row_counts, indent=2)
    )

    response = llm.chat.completions.create(
        messages=[{"role": "user", "content": final_prompt}],
        response_format={"type": "json_schema", "json_schema": {"name": "SyntheticDataset", "schema": SyntheticDataset.model_json_schema()}},
        temperature=0.7
    )

    choice = response.choices[0]
    response_content = choice.message.content

    if choice.finish_reason == "length":
        print(f"⚠️ WARNING: Chunk #{call_count} was truncated. Skipping.")
        continue
    if not response_content:
        print(f"⚠️ WARNING: Received empty content for chunk #{call_count}. Skipping.")
        continue

    try:
        chunk_data = json.loads(response_content)
        print(f"✅ Received and parsed chunk #{call_count}.")
        for table_name, rows in chunk_data.items():
            if table_name in all_synthetic_data and rows:
                all_synthetic_data[table_name].extend(rows)
        # Log progress
        for name, rows in all_synthetic_data.items():
             print(f"   - '{name}': {len(rows)} / {TOTAL_ROW_COUNTS[name]} rows")
    except json.JSONDecodeError as e:
        print(f"❌ ERROR: Failed to parse JSON for chunk #{call_count}. Reason: {e}. Skipping.")
    
    time.sleep(1)

# --- 4. Deduplicate and Write to DB ---
print("\n✨ Data generation complete. Aggregating, deduplicating, and saving to database...")

synthetic_data = all_synthetic_data
print("\n--- Deduplicating generated data ---")
for table_name, rows in synthetic_data.items():
    if not rows: continue
    initial_count = len(rows)
    df = pd.DataFrame(rows).drop_duplicates()
    final_count = len(df)
    synthetic_data[table_name] = df.to_dict('records')
    print(f" - Table '{table_name}': Removed {initial_count - final_count} duplicates ({initial_count} -> {final_count}).")

# Final trim to ensure exact counts
for table_name, total_rows_needed in TOTAL_ROW_COUNTS.items():
    if table_name in synthetic_data:
        synthetic_data[table_name] = synthetic_data[table_name][:total_rows_needed]

with duckdb.connect(SYNTHETIC_DB_PATH) as con:
    for table_name, rows in synthetic_data.items():
        if rows:
            df = pd.DataFrame(rows)
            schema_cols = schema_df[schema_df['name'] == table_name].iloc[0]['column_names']
            for col in schema_cols:
                if col not in df.columns: df[col] = None
            df = df[schema_cols]
            con.execute(f"CREATE OR REPLACE TABLE {table_name} AS SELECT * FROM df")
    
    print(f"\n✅ Synthetic training sandbox created at: {SYNTHETIC_DB_PATH}")
    print("Tables created:", con.sql("SHOW TABLES;").fetchall())

import duckdb
from tabulate import tabulate

SYNTHETIC_DB_PATH = "data/synthetic_openflights.db"

# Connect to the synthetic database
with duckdb.connect(SYNTHETIC_DB_PATH, read_only=True) as con:
    
    # Get the list of all tables created
    all_tables = [table[0] for table in con.sql("SHOW TABLES;").fetchall()]
    
    # Select the first 3 tables to display (or all if fewer than 3)
    tables_to_validate = all_tables[:3]

    print("--- Validating the first few tables in the synthetic sandbox ---\n")

    # Execute and print results for the selected tables
    for table_name in tables_to_validate:
        print(f"--- SELECT * FROM {table_name} LIMIT 3; ---")
        try:
            result_df = con.sql(f"SELECT * FROM {table_name} LIMIT 3;").df()
            if not result_df.empty:
                print(tabulate(result_df, headers='keys', tablefmt='psql'))
            else:
                print(f"(Table '{table_name}' is empty)")
        except Exception as e:
            print(f"Query failed for table '{table_name}': {e}")
        print("\n")

import pandas as pd
import json
import time
from pydantic import BaseModel, Field
from typing import List
from fireworks import LLM
import os
import duckdb
from dotenv import load_dotenv

load_dotenv()

# --- 1. Define Generation Parameters and Pydantic Model ---
llm = LLM(model="accounts/fireworks/models/qwen3-coder-480b-a35b-instruct", deployment_type="serverless", api_key=os.getenv("FIREWORKS_API_KEY"))  # Use Qwen3-coder for SQL queries
TOTAL_QUERIES_TO_GENERATE = 1000  # Note, some of these queries will likely be duplicates or invalid, reducing the final number used for fine-tuning
QUERIES_PER_API_CALL = 30

class SqlQueryBatch(BaseModel):
    queries: List[str] = Field(description=f"A list of exactly {QUERIES_PER_API_CALL} unique and diverse SQL queries.")

print(f"🎯 Goal: Generate {TOTAL_QUERIES_TO_GENERATE} unique queries in batches of {QUERIES_PER_API_CALL}.")

# --- 2. Get Clean Schema From Synthetic DB ---
with duckdb.connect(SYNTHETIC_DB_PATH, read_only=True) as con:
    schema_df = con.sql("DESCRIBE;").df()
    schema_for_prompt = schema_df.to_markdown(index=False)

# --- 3. Setup Base Prompt and Generation Loop ---
base_query_generation_prompt = f"""
You are an expert SQL data analyst. Your task is to generate unique and diverse SQL queries based on the database schema provided.
The queries should be realistic and cover a range of complexities and SQL features (JOINS, GROUP BY, aggregates, etc.).
Ensure you break ties with ORDER BY clauses so that the same queries produce the same results when executed against the database.
Write only the raw SQL query text and nothing else (i.e., no markdown formatting); your output should be a directly executable valid SQL query.
Make sure your queries do not return duplicate rows (i.e., GROUP BY all columns that are not aggregate functions).
Ensure the generated SQL is valid for DuckDB.

**Database Schema:**
{schema_for_prompt}
""".strip()

all_generated_queries = []
# Loop until we have enough queries
while len(all_generated_queries) < TOTAL_QUERIES_TO_GENERATE:
    print(f"\n📞 --- Generating batch #{len(all_generated_queries) // QUERIES_PER_API_CALL + 1} ---")

    # Create a summary of queries generated so far to prevent duplicates
    existing_queries_summary = ""
    if all_generated_queries:
        summary_parts = ["\nYou have already generated the following queries:\n"]
        for i, q in enumerate(all_generated_queries):
            summary_parts.append(f"{i+1}. {q}")
        existing_queries_summary = "\n".join(summary_parts)

    # Construct the final prompt for this iteration
    final_prompt = (
        base_query_generation_prompt +
        existing_queries_summary +
        f"\n\nNow, generate {QUERIES_PER_API_CALL} new, unique SQL queries, which cover different analytic scenarios and are not already in the list above. Return your response as a single JSON object adhering to the specified schema."
    )

    response = llm.chat.completions.create(
        messages=[{"role": "user", "content": final_prompt}],
        response_format={"type": "json_schema", "json_schema": {"name": "SqlQueryBatch", "schema": SqlQueryBatch.model_json_schema()}},
        temperature=0.8
    )

    response_content = response.choices[0].message.content
    if response_content:
        try:
            new_queries = json.loads(response_content).get("queries", [])
            all_generated_queries.extend(new_queries)
            print(f"   - Received {len(new_queries)} new queries. Total now: {len(all_generated_queries)} / {TOTAL_QUERIES_TO_GENERATE}")
        except json.JSONDecodeError as e:
            print(f"❌ ERROR: Failed to parse generated queries in this batch: {e}")
    
    time.sleep(1) # Be nice to the API

# --- 4. Deduplicate, Trim, and Save --- 
print("\n✨ Generation complete. Deduplicating and saving...")
initial_count = len(all_generated_queries)
# Simple, fast deduplication preserving order
unique_queries = list(dict.fromkeys(all_generated_queries))
final_count = len(unique_queries)
print(f" - Removed {initial_count - final_count} duplicates ({initial_count} -> {final_count}).")

# Trim to the exact number we need
final_queries = unique_queries[:TOTAL_QUERIES_TO_GENERATE]

# Save the final list to a file
QUERIES_FILE_PATH = "data/generated_queries.json"
with open(QUERIES_FILE_PATH, 'w') as f:
    json.dump({"queries": final_queries}, f, indent=2)

print(f"\n✅ Successfully saved {len(final_queries)} unique queries to `{QUERIES_FILE_PATH}`.")
print("\n--- Here are a few examples: ---")
for query in final_queries[:5]:
    print(f"- {query}")

# Check the proportion of generated SQL queries that return zero rows

import json
import duckdb

try:
    SYNTHETIC_DB_PATH
except NameError:
    SYNTHETIC_DB_PATH = "data/synthetic_openflights.db"

try:
    QUERIES_FILE_PATH
except NameError:
    QUERIES_FILE_PATH = "data/generated_queries.json"

with open(QUERIES_FILE_PATH, "r") as f:
    queries = json.load(f).get("queries", [])

total = len(queries)
zero_rows = 0
success = 0
failed = 0

with duckdb.connect(SYNTHETIC_DB_PATH, read_only=True) as con:
    for q in queries:
        try:
            cnt = con.sql(f"SELECT COUNT(*) AS c FROM ({q}) AS t").fetchone()[0]
            zero_rows += (cnt == 0)
            success += 1
        except Exception:
            failed += 1

print(f"Total queries: {total}")
print(f"Executed successfully: {success}")
print(f"Execution errors: {failed}")
print(f"Zero-row results: {zero_rows}")
if total > 0:
    print(f"Proportion zero-row (overall): {zero_rows/total:.3f}")
if success > 0:
    print(f"Proportion zero-row (successful only): {zero_rows/success:.3f}")

import os, json, re, time
import pandas as pd
import duckdb
from typing import List, Optional, Any, Dict, Type, Set
from pydantic import BaseModel, create_model
from fireworks import LLM
import datetime, decimal, uuid
from collections import defaultdict
from dotenv import load_dotenv

# --- Config ---
SYNTHETIC_DB_PATH = "data/synthetic_openflights.db"
QUERIES_FILE_PATH = "data/generated_queries.json"

# --- Setup LLM ---
load_dotenv()
llm = LLM(
    model="accounts/fireworks/models/qwen3-coder-480b-a35b-instruct",
    deployment_type="serverless",
    api_key=os.getenv("FIREWORKS_API_KEY"),
)

# --- Schema helpers ---
def map_sql_type_to_python(sql_type: str) -> Type:
    s = str(sql_type).upper()
    if "DECIMAL" in s: return decimal.Decimal
    if any(k in s for k in ("DOUBLE","FLOAT","REAL")): return float
    if any(k in s for k in ("BIGINT","INT")): return int
    if any(k in s for k in ("VARCHAR","TEXT","STRING")): return str
    if "TIMESTAMP" in s: return datetime.datetime
    if "DATE" in s: return datetime.date
    if "TIME" in s: return datetime.time
    if "BOOLEAN" in s: return bool
    if any(k in s for k in ("BLOB","BYTEA")): return bytes
    if "UUID" in s: return uuid.UUID
    return object

with duckdb.connect(SYNTHETIC_DB_PATH, read_only=True) as con_ro:
    schema_df = con_ro.sql("DESCRIBE;").df()
schema_for_prompt = schema_df.to_markdown(index=False)

def extract_tables_from_query(sql: str) -> Set[str]:
    """Extract table names from SQL query."""
    # Remove comments
    sql = re.sub(r'--.*$', '', sql, flags=re.MULTILINE)
    sql = re.sub(r'/\*.*?\*/', '', sql, flags=re.DOTALL)
    
    tables = set()
    # Match FROM and JOIN clauses
    patterns = [
        r'(?:FROM|JOIN)\s+([a-zA-Z_][a-zA-Z0-9_]*)',
        r'(?:FROM|JOIN)\s+"([^"]+)"',
        r'(?:FROM|JOIN)\s+`([^`]+)`',
    ]
    
    for pattern in patterns:
        matches = re.findall(pattern, sql, re.IGNORECASE)
        tables.update(matches)
    
    # Filter out SQL keywords that might be captured
    sql_keywords = {'select', 'where', 'group', 'order', 'having', 'limit', 'as', 'on', 'and', 'or', 'not', 'in', 'exists'}
    tables = {t for t in tables if t.lower() not in sql_keywords}
    
    return tables

def table_columns(name: str) -> List[str]:
    row = schema_df[schema_df["name"] == name]
    if row.empty:
        return []
    return list(row.iloc[0]["column_names"])

def table_types(name: str) -> List[str]:
    row = schema_df[schema_df["name"] == name]
    if row.empty:
        return []
    return list(row.iloc[0]["column_types"])

def build_rows_payload_model(tables: List[str]) -> Type[BaseModel]:
    per_table_row_models: Dict[str, Type[BaseModel]] = {}
    for t in tables:
        cols = table_columns(t)
        types = table_types(t)
        if not cols:
            continue
        fields: Dict[str, Any] = {}
        for c, ty in zip(cols, types):
            fields[c] = (Optional[map_sql_type_to_python(ty)], None)
        per_table_row_models[t] = create_model(f"{t.capitalize()}Row", **fields)
    
    payload_fields: Dict[str, Any] = {}
    for t, row_model in per_table_row_models.items():
        payload_fields[t] = (List[row_model], [])
    
    if not payload_fields:
        return create_model("RowsPayloadFallback", rows=(List[dict], []))
    return create_model("RowsPayload", **payload_fields)

def count_rows(con, sql: str) -> int:
    try:
        return con.sql(f"SELECT COUNT(*) AS c FROM ({sql}) AS t").fetchone()[0]
    except Exception:
        return -1

def get_sample_data(con, table: str, limit: int = 3) -> List[dict]:
    """Get sample rows from a table for context."""
    try:
        df = con.sql(f'SELECT * FROM "{table}" LIMIT {limit}').df()
        return df.to_dict("records")
    except Exception:
        return []

# --- Load queries ---
with open(QUERIES_FILE_PATH, "r") as f:
    queries = json.load(f).get("queries", [])

total_q = len(queries)
print(f"Total queries: {total_q}")

# --- Parameters ---
TARGET_MAX_ZERO_PERCENT = 10
MAX_ITERATIONS = 20
BATCH_SIZE = 10  # Process up to 10 queries at once
MAX_ROWS_PER_TABLE_PER_BATCH = 2  # Generate at most 2 rows per table per batch

# Initial assessment
with duckdb.connect(SYNTHETIC_DB_PATH, read_only=True) as con_check:
    zero_indices = [i for i, q in enumerate(queries) if count_rows(con_check, q) == 0]
    initial_zero_count = len(zero_indices)
    print(f"Initial zero-result queries: {initial_zero_count}/{total_q} ({initial_zero_count/total_q*100:.1f}%)")

# Track which queries have been attempted
processed_query_indices = set()
retry_count = 0  # Track how many times we've cycled through all queries

iteration = 0
with duckdb.connect(SYNTHETIC_DB_PATH) as con:
    while iteration < MAX_ITERATIONS:
        iteration += 1
        
        # Get ALL current zero-result queries
        all_zero_indices = [i for i, q in enumerate(queries) if count_rows(con, q) == 0]
        zero_count = len(all_zero_indices)
        zero_percent = (zero_count / total_q * 100) if total_q else 0
        
        print(f"\n[Iteration {iteration}] Zero-result: {zero_count}/{total_q} ({zero_percent:.1f}%)")
        
        if zero_percent <= TARGET_MAX_ZERO_PERCENT or zero_count == 0:
            print(f"✅ Target achieved! {zero_percent:.1f}% <= {TARGET_MAX_ZERO_PERCENT}%")
            break
        
        # Get unprocessed zero-result queries
        unprocessed_zero_indices = [i for i in all_zero_indices if i not in processed_query_indices]
        
        # If we've processed all zero-result queries, reset to try stubborn ones again
        if not unprocessed_zero_indices and all_zero_indices:
            retry_count += 1
            print(f"  All zero-result queries have been attempted. Starting retry cycle #{retry_count}")
            processed_query_indices.clear()
            unprocessed_zero_indices = all_zero_indices
            
            # If we've done multiple retry cycles with no progress, stop
            if retry_count > 2:
                print(f"  Stopping after {retry_count} retry cycles")
                break
        
        # Take a batch of unprocessed queries
        batch_indices = unprocessed_zero_indices[:BATCH_SIZE]
        if not batch_indices:
            print(f"[Iteration {iteration}] No queries to process.")
            break
            
        batch_queries = [queries[i] for i in batch_indices]
        processed_query_indices.update(batch_indices)
        
        print(f"  Processing batch: queries {batch_indices[:3]}{'...' if len(batch_indices) > 3 else ''} ({len(batch_indices)} total)")
        
        # Group queries by their involved tables for efficient processing
        query_tables_map = defaultdict(list)
        for idx, q in zip(batch_indices, batch_queries):
            tables = extract_tables_from_query(q)
            if tables:
                # Create a key from sorted table names
                key = tuple(sorted(tables))
                query_tables_map[key].append((idx, q))
        
        if not query_tables_map:
            print(f"  No tables found in batch. Moving to next batch.")
            continue
        
        total_fixed = 0
        
        # Process each group of queries that share the same tables
        for table_tuple, query_list in query_tables_map.items():
            tables = list(table_tuple)
            query_indices = [idx for idx, _ in query_list]
            query_texts = [q for _, q in query_list]
            
            print(f"    Processing {len(query_list)} queries involving tables: {tables}")
            
            # Build Pydantic model
            RowsPayload = build_rows_payload_model(tables)
            rows_schema = RowsPayload.model_json_schema()
            
            # Limit rows per table
            props = rows_schema.get("properties", {})
            for t in list(props.keys()):
                spec = props.get(t, {})
                if isinstance(spec, dict) and spec.get("type") == "array":
                    spec["maxItems"] = min(MAX_ROWS_PER_TABLE_PER_BATCH, len(query_list))
            
            # Get sample data for context
            existing_samples = {}
            for t in tables:
                samples = get_sample_data(con, t, limit=5)
                if samples:
                    existing_samples[t] = samples
            
            # Analyze query patterns to understand what's needed
            query_analysis = []
            for q in query_texts[:3]:  # Analyze first 3 queries for brevity
                query_analysis.append(f"- {q}")
            
            system_prompt = """You are an expert at generating natural, realistic database records.
Generate minimal new rows that will make the provided SQL queries return results.
The data should be diverse, realistic, and consistent with the domain implied by the schema."""
            
            user_prompt = f"""
Given this DuckDB schema and SQL queries that return zero rows, generate the minimum number
of natural, realistic rows needed to make these queries return results.

**Database Schema:**
{schema_for_prompt}

**Tables to populate:** {json.dumps(tables)}

**Example existing data (for style/format reference):**
{json.dumps(existing_samples, indent=2, default=str) if existing_samples else "No existing samples available"}

**Queries that need to return results ({len(query_texts)} total):**
{chr(10).join(query_analysis)}
{f"... and {len(query_texts) - 3} more similar queries" if len(query_texts) > 3 else ""}

**Requirements:**
1. Generate at most {MAX_ROWS_PER_TABLE_PER_BATCH} rows per table
2. Make the data satisfy the WHERE conditions and JOINs in the queries
3. Use natural, realistic values appropriate for the domain
4. Maintain referential integrity between tables
5. Avoid duplicating existing IDs - generate new unique IDs
6. Create diverse data that satisfies multiple queries if possible

Return ONLY the JSON object with the new rows. No explanations.
""".strip()
            
            try:
                resp = llm.chat.completions.create(
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt}
                    ],
                    response_format={
                        "type": "json_schema",
                        "json_schema": {"name": "RowsPayload", "schema": rows_schema}
                    },
                    temperature=0.5,  # Some randomness for natural data
                )
                
                content = resp.choices[0].message.content
                if not content:
                    print(f"      Empty response from LLM")
                    continue
                
                payload = json.loads(content)
                
                # Insert generated rows
                rows_inserted = 0
                for t in tables:
                    rows = payload.get(t, [])
                    if not rows or not isinstance(rows, list):
                        continue
                    
                    # Convert to DataFrame
                    df = pd.DataFrame(rows[:MAX_ROWS_PER_TABLE_PER_BATCH])
                    cols = table_columns(t)
                    if not cols:
                        continue
                    
                    # Align columns
                    for c in cols:
                        if c not in df.columns:
                            df[c] = None
                    df = df[cols]
                    
                    # Insert new rows (avoiding exact duplicates)
                    try:
                        con.register("new_rows_df", df)
                        con.execute(f'INSERT INTO "{t}" SELECT * FROM new_rows_df EXCEPT SELECT * FROM "{t}"')
                        con.unregister("new_rows_df")
                        rows_inserted += len(df)
                    except Exception as e:
                        print(f"      Warning: Failed to insert into {t}: {e}")
                
                # Check how many queries were fixed
                fixed_in_group = 0
                for idx in query_indices:
                    if count_rows(con, queries[idx]) > 0:
                        fixed_in_group += 1
                
                print(f"      Inserted {rows_inserted} rows, fixed {fixed_in_group}/{len(query_list)} queries")
                total_fixed += fixed_in_group
                
            except Exception as e:
                print(f"      Error processing group: {e}")
                continue
            
            time.sleep(0.5)  # Rate limiting
        
        print(f"  [Iteration {iteration}] Total fixed in this iteration: {total_fixed} queries")
        
        # If we're in a retry cycle and made no progress, stop
        if retry_count > 0 and total_fixed == 0:
            print(f"  No progress made in retry cycle. Stopping.")
            break

# Final cleanup - remove exact duplicates
print("\n--- Final deduplication ---")
with duckdb.connect(SYNTHETIC_DB_PATH) as con:
    tables = con.sql("""
        SELECT table_name 
        FROM information_schema.tables 
        WHERE table_type = 'BASE TABLE' 
        AND table_schema = 'main'
    """).fetchall()
    
    for (table_name,) in tables:
        before = con.sql(f'SELECT COUNT(*) FROM "{table_name}"').fetchone()[0]
        con.execute(f'CREATE OR REPLACE TABLE "{table_name}" AS SELECT DISTINCT * FROM "{table_name}"')
        after = con.sql(f'SELECT COUNT(*) FROM "{table_name}"').fetchone()[0]
        if before != after:
            print(f"  {table_name}: removed {before - after} duplicate rows")

# Final report
with duckdb.connect(SYNTHETIC_DB_PATH, read_only=True) as con_read:
    final_zero_indices = [i for i, q in enumerate(queries) if count_rows(con_read, q) == 0]
    final_zero_count = len(final_zero_indices)
    final_zero_percent = (final_zero_count / total_q * 100) if total_q else 0
    
    print(f"\n{'='*60}")
    print(f"AUGMENTATION COMPLETE")
    print(f"{'='*60}")
    print(f"Initial zero-result: {initial_zero_count}/{total_q} ({initial_zero_count/total_q*100:.1f}%)")
    print(f"Final zero-result:   {final_zero_count}/{total_q} ({final_zero_percent:.1f}%)")
    print(f"Improvement:         {initial_zero_count - final_zero_count} queries now return results")
    
    if final_zero_percent <= TARGET_MAX_ZERO_PERCENT:
        print(f"\n✅ SUCCESS: Achieved target of ≤{TARGET_MAX_ZERO_PERCENT}% zero-result queries")
    else:
        print(f"\n⚠️  Partial success. Consider running again or adjusting parameters.")
        
    # Show a few examples of remaining zero-result queries if any
    if final_zero_count > 0 and final_zero_count <= 5:
        print(f"\nRemaining zero-result queries:")
        for idx in final_zero_indices[:5]:
            print(f"  [{idx}]: {queries[idx][:100]}...")

import duckdb
import json
import pandas as pd

# --- 1. Define File Paths ---
SYNTHETIC_DB_PATH = "data/synthetic_openflights.db"
QUERIES_FILE_PATH = "data/generated_queries.json"
GROUND_TRUTH_FILE_PATH = "data/ground_truth_results.jsonl"

# Thresholds to drop overly large results
MAX_RESULT_ROWS = 1000          # skip if result has more than this many rows
MAX_RESULT_BYTES = 100_000      # ~0.1MB; skip if JSON payload exceeds this

# --- 2. Load Generated Queries ---
with open(QUERIES_FILE_PATH, 'r') as f:
    queries_data = json.load(f)
    queries_to_execute = queries_data.get("queries", [])

print(f"Loaded {len(queries_to_execute)} queries to execute.")

# --- 3. Execute Queries and Store Results ---
ground_truth_results = []
successful_executions = 0
failed_executions = 0
oversized_skipped = 0

print("Executing queries against the synthetic database...")
with duckdb.connect(SYNTHETIC_DB_PATH, read_only=True) as con:
    for query in queries_to_execute:
        try:
            result_df = con.sql(query).df()

            # Replace any NaN/NaT values with None, which serializes to JSON `null`
            result_df = result_df.astype(object).where(pd.notna(result_df), None)
            result_records = result_df.to_dict('records')

            # Skip examples that are "much too wide"
            if len(result_records) > MAX_RESULT_ROWS:
                oversized_skipped += 1
                continue
            # Size-based guard (in bytes)
            payload_bytes = len(json.dumps(result_records, ensure_ascii=False).encode("utf-8"))
            if payload_bytes > MAX_RESULT_BYTES:
                oversized_skipped += 1
                continue

            ground_truth_results.append({
                "query": query,
                "result": result_records
            })
            successful_executions += 1
        except Exception as e:
            print(f"⚠️  Skipping query due to execution error: {query}\n   Error: {e}\n")
            failed_executions += 1

print(f"\nExecution complete. Success: {successful_executions}, Skipped (oversized): {oversized_skipped}, Failed: {failed_executions}.")

# --- 4. Save the Ground-Truth Data ---
with open(GROUND_TRUTH_FILE_PATH, 'w') as f:
    for entry in ground_truth_results:
        f.write(json.dumps(entry) + '\n')

print(f"\n✅ Successfully saved {len(ground_truth_results)} ground-truth results to `{GROUND_TRUTH_FILE_PATH}`.")

# --- 5. Print an Example ---
if ground_truth_results:
    print("\n--- Example ground_truth_results dataset entry ---")
    print(json.dumps(ground_truth_results[0], indent=2))

import json
import time
import jsonlines
from typing import List
import random
from fireworks import LLM
import os

# --- 1. Define File Paths and Parameters ---
llm = LLM(model="accounts/fireworks/models/qwen3-coder-480b-a35b-instruct", deployment_type="serverless", api_key=os.getenv("FIREWORKS_API_KEY"))
GROUND_TRUTH_FILE_PATH = "data/ground_truth_results.jsonl"
FINAL_TRAINING_DATA_PATH = "data/final_rft_sql_train_data.jsonl"
FINAL_TEST_DATA_PATH = "data/final_rft_sql_test_data.jsonl"

# --- 2. Load Ground-Truth Data ---
query_result_pairs = []
with jsonlines.open(GROUND_TRUTH_FILE_PATH) as reader:
    for obj in reader:
        query_result_pairs.append(obj)

print(f"Loaded {len(query_result_pairs)} query-result pairs.")

# --- 3. Use LLM to Generate Natural Language Questions ---
nl_generation_prompt_template = """
You are an expert data analyst who is great at translating SQL queries into plain English.
Based on the database schema and the provided SQL query, what is a natural language question a business user would ask to get this information?
Ensure that the question is precise enough to accurately map to the corresponding SQL query.

**Database Schema:**
{schema_for_prompt}

**SQL Query:**
{query}

Provide only the user's question, without any preamble or explanation.
""".strip()

# The system prompt that will be included in the final training data for the RFT job.
# It gives the model its instructions at inference time.
rft_system_prompt = f"""
You are an expert SQL data analyst.
Your task is to write a single, valid DuckDB SQL query to answer the user's question, based on the provided database schema.
Write only the raw SQL query text and nothing else (i.e., no markdown formatting); your output should be a directly executable valid SQL query.
Make sure your queries do not return duplicate rows (i.e., GROUP BY all columns that are not aggregate functions).
Ensure the generated SQL is valid for DuckDB.

**Database Schema:**
{schema_for_prompt}
""".strip()

final_generated_data = []
print(f"Generating natural language questions and formatting for RFT for {len(query_result_pairs)} queries...")

for i, pair in enumerate(query_result_pairs):
    print(f" - Processing query {i+1}/{len(query_result_pairs)}...")
    query = pair['query']
    ground_truth = pair['result']
    nl_generation_prompt = nl_generation_prompt_template.format(schema_for_prompt=schema_for_prompt, query=query)
    
    response = llm.chat.completions.create(
        messages=[{"role": "user", "content": nl_generation_prompt}],
        temperature=0.5
    )
    
    nl_question = response.choices[0].message.content
    if nl_question:  # Only include the entry if the LLM generated a question
        # Assemble the final data structure
        rft_entry = {
            "messages": [
                {"role": "system", "content": rft_system_prompt},
                {"role": "user", "content": nl_question.strip()},
                {"role": "assistant", "content": query}
            ],
            "ground_truth": ground_truth  # The ground-truth result for the evaluator
        }
        final_generated_data.append(rft_entry)
    
    time.sleep(0.5) # Be nice to the API

# --- 4. Shuffle and Split the Dataset ---
print(f"\nGenerated {len(final_generated_data)} total examples.")

# Filter out entries where 'ground_truth' is an empty list
original_count = len(final_generated_data)
final_generated_data = [entry for entry in final_generated_data if entry.get("ground_truth")]
print(f"Filtered out {original_count - len(final_generated_data)} examples with empty ground truth.")

print(f"Now splitting the remaining {len(final_generated_data)} examples into train and test sets.")

random.seed(42)
random.shuffle(final_generated_data)

split_index = int(len(final_generated_data) * 0.8)
train_data = final_generated_data[:split_index]
test_data = final_generated_data[split_index:]

print(f"Train set size: {len(train_data)}")
print(f"Test set size: {len(test_data)}")

# --- 5. Save the Final RFT-Ready Datasets ---
with jsonlines.open(FINAL_TRAINING_DATA_PATH, mode='w') as writer:
    writer.write_all(train_data)
print(f"\n✅ Successfully saved training dataset to `{FINAL_TRAINING_DATA_PATH}`.")

with jsonlines.open(FINAL_TEST_DATA_PATH, mode='w') as writer:
    writer.write_all(test_data)
print(f"✅ Successfully saved test dataset to `{FINAL_TEST_DATA_PATH}`.")

# --- 6. Print an Example ---
if train_data:
    print("\n--- Example RFT training entry ---")
    print(json.dumps(train_data[0], indent=2))

    import os, contextlib, uvicorn
    from starlette.applications import Starlette
    from starlette.routing import Mount
    from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
    from mcp_server_motherduck import build_application

    DB = "data/synthetic_openflights.db"          # ← path from previous steps
    PORT = int(os.environ.get("PORT", 8080))        # Cloud Run injects $PORT

    # 1️⃣ Build the core SQL-aware MCP server (read-only for safety).
    server, _ = build_application(db_path=DB, read_only=True)

    # 2️⃣ Wrap it so HTTP clients can talk to it (ASGI handler).
    sess = StreamableHTTPSessionManager(app=server, event_store=None, stateless=True)

    async def handler(scope, receive, send):
        await sess.handle_request(scope, receive, send)

    @contextlib.asynccontextmanager
    async def lifespan(app):
        async with sess.run():
            yield                                        # keep sessions alive

    # 3️⃣ Starlette turns that handler into a full ASGI app Uvicorn can serve.
    app = Starlette(routes=[Mount("/mcp", app=handler)], lifespan=lifespan)

    if __name__ == "__main__":
        print(f"🔥 MCP endpoint → http://0.0.0.0:{PORT}/mcp")
        uvicorn.run(app, host="0.0.0.0", port=PORT)

mcp
mcp-server-motherduck
duckdb
uvicorn
starlette

base
FROM python:3.11-slim
WORKDIR /app

COPY mcp_requirements.txt .
RUN pip install --no-cache-dir -r mcp_requirements.txt

COPY run_mcp_server.py .
COPY data/synthetic_openflights.db ./data/

EXPOSE 8080

CMD ["python", "run_mcp_server.py"]

# .gcloudignore

# 1. Ignore EVERYTHING in the directory by default.
*

# 2. Now, create exceptions for ONLY the files needed by the Dockerfile.
# The "!" character means "do not ignore this file".

# The Dockerfile itself is needed for the build process.
!Dockerfile

# The files explicitly copied by your Dockerfile:
!mcp_requirements.txt
!run_mcp_server.py

# 3. To include a specific file in a subdirectory, use this
#    three-line pattern to un-ignore the directory, re-ignore its
#    contents, and then un-ignore the specific file.
!data/
data/*
!data/synthetic_openflights.db

gcloud run deploy mcp-sql-rft-server \
  --source . \
  --region < YOUR_GCP_REGION > \
  --project < YOUR_GCP_PROJECT_ID > \
  --allow-unauthenticated \
  --port 8080

gcloud run services describe mcp-sql-rft-server \
--project < YOUR_GCP_PROJECT_ID > \
--region < YOUR_GCP_REGION > \
--format="value(status.url)"

curl -X POST "< YOUR_MCP_SERVER_URL_FROM_STEP_i >/mcp/" \
-H "Content-Type: application/json" \
-H "Accept: application/json, text/event-stream" \
-d '{
    "id": "list-tools-1",
    "jsonrpc": "2.0",
    "method": "tools/list",
    "params": {
        "session": {"id": "test-from-my-laptop"}
    }
}'

curl -X POST "< YOUR_MCP_SERVER_URL_FROM_STEP_i >/mcp/" \
-H "Content-Type: application/json" \
-H "Accept: application/json, text/event-stream" \
-d '{
    "id": "query-1",
    "jsonrpc": "2.0",
    "method": "tools/call",
    "params": {
        "session": {"id": "test-from-my-laptop"},
        "name": "query",
        "arguments": {
            "query": "SELECT COUNT(*) FROM airlines;"
        }
    }
}'

import requests
import json
import math

MCP_SERVER_URL = None  # <--- PUT MCP SERVER URL HERE without the /mcp/ suffix at the end

def evaluate(messages: list[dict], ground_truth: list[dict], **kwargs) -> dict:
    """
    Evaluates the model's generated SQL query by executing it against a live
    MCP server and comparing the result with the ground_truth.
    """
    
    def parse_duckdb_ascii_table(table_string: str) -> list[dict]:
        """
        Parses a DuckDB-style ASCII table string into a list of dictionaries.
        This version robustly handles 'NULL' values and empty strings.
        """
        lines = table_string.strip().split('\n')
        content_lines = [line for line in lines if line.strip() and not line.startswith('+')]
        if len(content_lines) < 2:
            return []
        
        header_raw = [h.strip() for h in content_lines[0].split('|')[1:-1]]
        data_lines = content_lines[1:]
        
        if len(data_lines) > 0:
            try:
                first_data_values = [v.strip() for v in data_lines[0].split('|')[1:-1]]
                if len(first_data_values) == len(header_raw) and all(v.isupper() for v in first_data_values):
                    data_lines = data_lines[1:]
            except IndexError:
                pass

        rows = []
        for line in data_lines:
            try:
                values_raw = [v.strip() for v in line.split('|')[1:-1]]
                if len(values_raw) == len(header_raw):
                    row_dict = {}
                    for i, header in enumerate(header_raw):
                        value_str = values_raw[i]
                        if value_str.upper() == 'NULL' or value_str == '':
                            row_dict[header] = None
                            continue
                        
                        try:
                            if '.' in value_str:
                                row_dict[header] = float(value_str)
                            else:
                                row_dict[header] = int(value_str)
                        except (ValueError, TypeError):
                            row_dict[header] = value_str
                    rows.append(row_dict)
            except IndexError:
                continue
        return rows

    # --- 1. Get MCP Server URL from Environment Variables ---
    mcp_server_url = MCP_SERVER_URL
    if not mcp_server_url:
        return {"score": 0, "is_score_valid": False, "reason": "FATAL: MCP_SERVER_URL environment variable is not set."}

    # --- 2. Get the SQL query from the model's response ---
    sql_query = messages[-1]['content'].strip()
    if not sql_query:
        return {"score": 0, "reason": "Model returned an empty response."}

    # --- 3. Execute the Query against the MCP Server ---
    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json, text/event-stream"
    }
    payload = {
        "id": "eval-query-1", "jsonrpc": "2.0", "method": "tools/call",
        "params": {"session": {"id": "stateless-eval-session"}, "name": "query", "arguments": {"query": sql_query}}
    }
    try:
        with requests.post(f"{mcp_server_url}/mcp/", headers=headers, json=payload, timeout=15, stream=True) as response:
            response.raise_for_status()
            response_data = None
            for line in response.iter_lines():
                if line:
                    decoded_line = line.decode('utf-8')
                    if decoded_line.startswith('data:'):
                        json_part = decoded_line[len('data:'):].strip()
                        if json_part:
                            response_data = json.loads(json_part)
                            break
            if response_data is None:
                return {"score": 0, "reason": "Could not find JSON data in event stream response from MCP server."}

        if "error" in response_data:
            return {"score": 0, "reason": f"SQL execution failed. Error: {response_data['error'].get('message', 'Unknown')}"}

        ascii_table = response_data['result']['content'][0]['text']
        predicted_rows = parse_duckdb_ascii_table(ascii_table)

    except requests.exceptions.RequestException as e:
        return {"score": 0, "reason": f"Network error calling MCP server: {e}"}
    except json.JSONDecodeError as e:
        return {"score": 0, "reason": f"JSON decode error from server response: {e}"}
    except (KeyError, IndexError):
        return {"score": 0, "reason": f"Failed to parse predicted result from MCP server response structure. Data found: {json.dumps(response_data)}"}
    except Exception as e:
        return {"score": 0, "reason": f"An unexpected error occurred during query execution: {e}"}

    # --- 4. Process Ground Truth ---
    if not isinstance(ground_truth, list):
        return {"score": 0, "is_score_valid": False, "reason": f"FATAL: ground_truth is not a list as expected. Got type: {type(ground_truth)}"}
    ground_truth_rows = ground_truth


    # --- 5. Comparison Logic ---
    def normalize_and_stringify(v):
        """
        Normalizes numbers and None before string conversion.
        """
        if v is None:
            return str(v)
        
        if isinstance(v, float) and not math.isinf(v) and not math.isnan(v) and v == int(v):
            v = int(v)
        return str(v)

    try:
        gt_values = sorted([sorted(map(normalize_and_stringify, row.values())) for row in ground_truth_rows])
        predicted_values = sorted([sorted(map(normalize_and_stringify, row.values())) for row in predicted_rows])

        if gt_values == predicted_values:
            score = 1
            reason = "Success: The SQL query produced the exact expected result."
        else:
            score = 0
            gt_json = json.dumps(ground_truth_rows)
            pred_json = json.dumps(predicted_rows)
            reason = f"Incorrect result. Expected (from ground_truth): {gt_json}. Got (from query): {pred_json}."
    
    except Exception as e:
        return {"score": 0, "reason": f"Error during result comparison: {e}"}

    return {"score": score, "reason": reason}

import requests
import json
import os
from fireworks import LLM

# --- 1. SETUP: Define API keys, server URLs, and the model to use ---

# IMPORTANT: Make sure your FIREWORKS_API_KEY is set as an environment variable.
# You can get one from https://fireworks.ai
if "FIREWORKS_API_KEY" not in os.environ:
    print("FATAL: FIREWORKS_API_KEY environment variable not set.")
    # If not set, you can hardcode it here for testing, but this is not recommended:
    # os.environ["FIREWORKS_API_KEY"] = "YOUR_API_KEY_HERE"

# The model we'll use to generate the SQL. This acts as our "base" model.
LLM_MODEL = "accounts/fireworks/models/llama-v3p1-8b-instruct"
llm = LLM(model=LLM_MODEL, deployment_type="auto", api_key=os.getenv("FIREWORKS_API_KEY"))

# The URL for your running MCP server.
MCP_SERVER_URL = None  # PUT MCP SERVER URL HERE without the /mcp/ suffix at the end


# --- 2. LOAD THE EXAMPLE DATA ---

# This is the example data you provided.
DATASET_FILE_PATH = "data/final_rft_sql_train_data.jsonl"
ROW_INDEX_TO_TEST = 0  # 0 is the first row, 1 is the second row, etc.

EXAMPLE_DATA = None
try:
    with open(DATASET_FILE_PATH, 'r') as f:
        for i, line in enumerate(f):
            if i == ROW_INDEX_TO_TEST:
                EXAMPLE_DATA = json.loads(line)
                break
    
    if EXAMPLE_DATA is None:
        with open(DATASET_FILE_PATH, 'r') as f:
            line_count = sum(1 for line in f)
        raise IndexError(f"row index {ROW_INDEX_TO_TEST} is out of bounds for file with {line_count} rows.")

    print(f"Successfully loaded row {ROW_INDEX_TO_TEST} from '{DATASET_FILE_PATH}'.\n")
    print(EXAMPLE_DATA)
    print()

except Exception as e:
    print(f"Warning: Could not load from file. Reason: {e}")

# If loading from file failed for any reason, use the hardcoded fallback data.
if EXAMPLE_DATA is None:
    print("Using hardcoded fallback EXAMPLE_DATA.\n")
    EXAMPLE_DATA = {
        "messages": [
            {"role": "system", "content": "\nYou are an expert SQL data analyst. Your task is to write a single, valid DuckDB SQL query to answer the user's question, based on the provided database schema. Do not provide any explanation or text other than the SQL query itself.\n\n**Database Schema:**\n| database              | schema   | name      | column_names                                                               | column_types                                                          | temporary   |\n|:----------------------|:---------|:----------|:---------------------------------------------------------------------------|:----------------------------------------------------------------------|:------------|\n| synthetic_openflights | main     | airlines  | ['airline_id' 'name' 'alias' 'iata' 'icao' 'callsign' 'country' 'active']  | ['BIGINT' 'VARCHAR' 'VARCHAR' 'VARCHAR' 'VARCHAR' 'VARCHAR' 'VARCHAR' | False       |\n|                       |          |           |                                                                            |  'VARCHAR']                                                           |             |\n| synthetic_openflights | main     | airports  | ['airport_id' 'name' 'city' 'country' 'iata' 'icao' 'latitude' 'longitude' | ['BIGINT' 'VARCHAR' 'VARCHAR' 'VARCHAR' 'VARCHAR' 'VARCHAR' 'DOUBLE'  | False       |\n|                       |          |           |  'altitude' 'timezone' 'dst' 'tz_db' 'type' 'source']                      |  'DOUBLE' 'BIGINT' 'DOUBLE' 'VARCHAR' 'VARCHAR' 'VARCHAR' 'VARCHAR']  |             |\n| synthetic_openflights | main     | countries | ['name' 'iso_code' 'dafif_code']                                           | ['VARCHAR' 'VARCHAR' 'VARCHAR']                                       | False       |\n| synthetic_openflights | main     | planes    | ['name' 'iata' 'icao']                                                     | ['VARCHAR' 'VARCHAR' 'VARCHAR']                                       | False       |\n| synthetic_openflights | main     | routes    | ['airline' 'airline_id' 'source_airport' 'source_airport_id'               | ['VARCHAR' 'BIGINT' 'VARCHAR' 'BIGINT' 'VARCHAR' 'BIGINT' 'VARCHAR'   | False       |\n|                       |          |           |  'destination_airport' 'destination_airport_id' 'codeshare' 'stops'        |  'BIGINT' 'VARCHAR']                                                  |             |\n|                       |          |           |  'equipment']                                                              |                                                                       |             |\n"},
            {"role": "user", "content": "Which countries have the most airlines, and how many airlines are there in each country, listed in descending order by the number of airlines and then alphabetically by country name?"},
            {"role": "assistant", "content": "SELECT country, COUNT(*) AS airline_count FROM airlines GROUP BY country ORDER BY airline_count DESC, country ASC"}
        ],
        "ground_truth": [{"country": "Canada", "airline_count": 10}, {"country": "Sweden", "airline_count": 10}, {"country": "Kenya", "airline_count": 9}, {"country": "United States", "airline_count": 9}, {"country": "Australia", "airline_count": 8}, {"country": "Spain", "airline_count": 6}, {"country": "Italy", "airline_count": 4}, {"country": "Switzerland", "airline_count": 4}, {"country": "Finland", "airline_count": 3}, {"country": "France", "airline_count": 3}, {"country": "Mexico", "airline_count": 3}, {"country": "Costa Rica", "airline_count": 2}, {"country": "Germany", "airline_count": 2}, {"country": "Iceland", "airline_count": 2}, {"country": "Ireland", "airline_count": 2}, {"country": "Japan", "airline_count": 2}, {"country": "Norway", "airline_count": 2}, {"country": "Singapore", "airline_count": 2}, {"country": "United Kingdom", "airline_count": 2}, {"country": "Argentina", "airline_count": 1}, {"country": "Brazil", "airline_count": 1}, {"country": "China", "airline_count": 1}, {"country": "Egypt", "airline_count": 1}, {"country": "Fiji", "airline_count": 1}, {"country": "Greece", "airline_count": 1}, {"country": "India", "airline_count": 1}, {"country": "Jordan", "airline_count": 1}, {"country": "Netherlands", "airline_count": 1}, {"country": "New Zealand", "airline_count": 1}, {"country": "Portugal", "airline_count": 1}, {"country": "Saudi Arabia", "airline_count": 1}, {"country": "South Africa", "airline_count": 1}, {"country": "Thailand", "airline_count": 1}, {"country": "United Arab Emirates", "airline_count": 1}]
    }

# Extract the prompts and ground truth from the data
system_prompt = EXAMPLE_DATA["messages"][0]["content"]
user_prompt = EXAMPLE_DATA["messages"][1]["content"]
GROUND_TRUTH_ROWS = EXAMPLE_DATA["ground_truth"]

# --- 3. HELPER FUNCTION: To parse the server's ASCII table response ---

def parse_duckdb_ascii_table(table_string: str) -> list[dict]:
    """
    Parses a DuckDB-style ASCII table string into a list of dictionaries.
    This version robustly handles 'NULL' values and empty strings.
    """
    lines = table_string.strip().split('\n')
    content_lines = [line for line in lines if line.strip() and not line.startswith('+')]
    if len(content_lines) < 2:
        return []
    
    header_raw = [h.strip() for h in content_lines[0].split('|')[1:-1]]
    data_lines = content_lines[1:]
    
    if len(data_lines) > 0:
        try:
            first_data_values = [v.strip() for v in data_lines[0].split('|')[1:-1]]
            if len(first_data_values) == len(header_raw) and all(v.isupper() for v in first_data_values):
                data_lines = data_lines[1:]
        except IndexError:
            pass

    rows = []
    for line in data_lines:
        try:
            values_raw = [v.strip() for v in line.split('|')[1:-1]]
            if len(values_raw) == len(header_raw):
                row_dict = {}
                for i, header in enumerate(header_raw):
                    value_str = values_raw[i]
                    if value_str.upper() == 'NULL' or value_str == '':
                        row_dict[header] = None
                        continue
                    
                    try:
                        if '.' in value_str:
                            row_dict[header] = float(value_str)
                        else:
                            row_dict[header] = int(value_str)
                    except (ValueError, TypeError):
                        row_dict[header] = value_str
                rows.append(row_dict)
        except IndexError:
            continue
    return rows

# --- 4. GENERATE SQL QUERY USING THE LLM ---

print("="*20)
print("LLM QUERY GENERATION")
print("="*20)

model_generated_sql = ""
try:
    print(f"User prompt: {user_prompt}")
    print(f"Ground truth: {GROUND_TRUTH_ROWS}")
    print(f"Calling model '{LLM_MODEL}' to generate SQL query...")
    
    messages_for_llm = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]
    
    response = llm.chat.completions.create(
        model=LLM_MODEL,
        messages=messages_for_llm,
        temperature=0.0  # Set to 0 for deterministic output
    )
    
    model_generated_sql = response.choices[0].message.content.strip()
    print("\nModel Generated SQL Query:")
    print(model_generated_sql)
    
except Exception as e:
    print(f"\nAN ERROR OCCURRED during LLM call: {e}")


# --- 5. EXECUTE GENERATED QUERY ON MCP SERVER ---

predicted_rows = []
if model_generated_sql:
    try:
        print("\n" + "="*20)
        print("MCP SERVER EXECUTION")
        print("="*20)
        print(f"Sending query to MCP server...")
        
        headers = {"Content-Type": "application/json", "Accept": "application/json, text/event-stream"}
        payload = {
            "id": "eval-query-1", "jsonrpc": "2.0", "method": "tools/call",
            "params": {"session": {"id": "stateless-eval-session"}, "name": "query", "arguments": {"query": model_generated_sql}}
        }

        with requests.post(f"{MCP_SERVER_URL}/mcp/", headers=headers, json=payload, timeout=20, stream=True) as response:
            response.raise_for_status()
            response_data = None
            for line in response.iter_lines():
                if line and line.decode('utf-8').startswith('data:'):
                    json_part = line.decode('utf-8')[len('data:'):].strip()
                    if json_part:
                        response_data = json.loads(json_part)
                        break
            
            if response_data is None: raise RuntimeError("No JSON data in event stream.")
            if "error" in response_data: raise RuntimeError(f"SQL Error: {response_data['error'].get('message', 'Unknown')}")

            ascii_table = response_data['result']['content'][0]['text']
            predicted_rows = parse_duckdb_ascii_table(ascii_table)
            print("\nParsed Result from Server:")
            print(json.dumps(predicted_rows, indent=2))

    except Exception as e:
        print(f"\nAN ERROR OCCURRED during MCP call: {e}")

# --- 6. FINAL COMPARISON ---
print("\n" + "="*20)
print("COMPARISON")
print("="*20)

if not predicted_rows:
    print("Skipping comparison: no rows returned from query or an error occurred.")
else:
    gt_values = sorted([sorted(map(str, row.values())) for row in GROUND_TRUTH_ROWS])
    predicted_values = sorted([sorted(map(str, row.values())) for row in predicted_rows])

    if gt_values == predicted_values:
        print("\n✅ GOOD RESULT: The base model generated SQL that produced the correct data.\n")
    else:
        print("\n❌ BAD RESULT: The base model's SQL produced different data than expected.\n")
        print("This is often the intended outcome when testing a base model, as it highlights what fine-tuning needs to correct.")