import pykx as kx
# Load a sample CSV and push it into q as 'trades'
pandas_tbl = pd.read_csv(
    'https://huggingface.co/datasets/opensporks/stocks/resolve/main/stocks.csv',
    parse_dates=['timestamp']
)
tbl = kx.toq(pandas_tbl)        # DataFrame → q Table
tbl = tbl.rename(columns={"close": "closePrice"})
kx.q['trades'] = tbl             # assign to q global

# get table description and save as text stored in a variable:
table_desc = tbl.dtypes
print(table_desc)

# get min and last date
min_date = kx.q('first trades`timestamp').pd()
max_date = kx.q('last trades`timestamp').pd()

print(f"\nMinimum timestamp: {min_date}")
print(f"Maximum timestamp: {max_date}")

SYSTEM_PROMPT = 'Create a valid sql query based on this table schema and natural language query'

class TranslateResponse(BaseModel):
    qsql: str

def translate_to_qsql(nl_query: str, table: str, schema: str) -> str:
    system = (
        SYSTEM_PROMPT
    )
    user = (
        f"Table: {table}\n"
        f"Schema: {schema}\n"
        f"Question: {nl_query}"
    )
    resp = client.beta.chat.completions.parse(
        model="o3",
        messages=[
            {"role": "system", "content": system},
            {"role": "user",   "content": user},
        ],
        response_format=TranslateResponse,
    )
    return resp.parsed.qsql

def visualize_df_with_plotly(df, model="gpt-4o", sample_size=20, temperature=0):
    """
    Given a pandas DataFrame df and an OpenAI-compatible client,
    infers a schema and sample, asks the model to generate Plotly Express code,
    and then executes and displays the resulting chart.
    """
    # 1. Build schema and sample, converting timestamps to strings
    sample_df = df.head(sample_size).copy()
    for col, dtype in sample_df.dtypes.items():
        if "datetime" in str(dtype).lower() or "timestamp" in str(dtype).lower():
            sample_df[col] = sample_df[col].astype(str)
    sample = sample_df.to_dict(orient="list")
    schema_auto = {col: str(dtype) for col, dtype in df.dtypes.items()}

# 2. Prepare prompts
    system = (
        "You're an expert Python data-visualizer. "
        "Given a DataFrame schema and sample data, "
        "write self-contained code using Plotly Express that "
        "takes the df and, creates a fig named fig, "
        "Write code that assigns df to fig, "
        "but do not call fig.show(). I will handle displaying. "
        "Feel free to show any other aggregations and print as numbers as you feel is needed as long as the figure is made"
        "GIVE ME ONLY CODE AND NOTHING ELSE."
        "Do not rewrite the values in the df, assume it is global in a dataframe called df"
        "Think about what the user wants in the query and provide it"
        "Never hardcode values"
    )
    user = (
        f"Schema: {json.dumps(schema_auto)}\n\n"
        f"Sample data: {json.dumps(sample, default=str)}"
    )
    # 3. Call the chat completion API
    resp = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system},
            {"role": "user",   "content": user}
        ],
        temperature=temperature,
    )
    # 4. Extract the code from the response
    chart_code = resp.choices[0].message.content
    if chart_code.startswith("```"):
        chart_code = chart_code.split("\n", 1)[1].rsplit("\n```", 1)[0]
    print(chart_code)
    # 5. Execute the generated code
    local_vars = {"df": df, "px": px}
    exec(chart_code, {}, local_vars)
    # 6. Display the figure
    fig = local_vars.get("fig")
    if fig:
        fig.show()
    else:
        raise RuntimeError("The generated code did not produce a fig object.")

def query_and_plot(nl_query: str, table: str = "trades") -> None:
    """
    Given a natural-language question about the kdb+ table,
    translates it to q-SQL, runs it, and renders an appropriate chart.
    """
    # --- translate to q-SQL ---
    # derive schema from the table in q
    q_tbl = tbl
    schema = str(q_tbl.dtypes)
    print(schema)

    # call the existing translator (assumes translate_to_qsql is defined globally)
    qsql = translate_to_qsql(nl_query, table, schema)
    print("🔧 Generated q-SQL:\n", qsql)
    # --- execute query ---
    result = kx.q.sql(qsql)
    df = result.pd()
    print("Resulting df After AI Query")
    display(df.head())
    print("\n── Resulting df Description ──")
    display(df.describe(include="all"))
    visualize_df_with_plotly(df)

query_and_plot("Chart the cumulative returns for TSLA and AAPL")

query_and_plot("what is the average price for each stock?"

query_and_plot("Which day had the highest closing price for each stock?")