import json
from langchain.chat_models import ChatOpenAI
from langchain.schema import SystemMessage, HumanMessage

def generate_sql_and_chart_details(schema_str, nlq_query, openai_api_key):
    
    client = ChatOpenAI(model="gpt-4.1", openai_api_key=openai_api_key)
    
# This will be our new prompt template
    SYSTEM_PROMPT = "You are an expert AI assistant specializing in converting natural language business queries into PostgreSQL-compliant SQL and corresponding JSON chart specifications. You must follow all instructions precisely and output only a single, valid JSON object."

    HUMAN_PROMPT_TEMPLATE = """
### TASK
Your task is to analyze the provided database schema and the user's natural language query (NLQ). Based on this analysis, you will generate a PostgreSQL query to retrieve the necessary data and a JSON specification for a suitable chart to visualize the results.

### DATABASE SCHEMA
Here are the tables and columns you must use. Adhere strictly to these names.

{schema}

### DOMAIN-SPECIFIC KNOWLEDGE
Pay close attention to these business definitions when constructing your query:

- **Status Values**: The 'status' column can only be 'Stopped' or 'Running'.
- **No-Metal Time**: This is a key metric calculated as the total time (in minutes or hours) a machine was stopped due to the material level being too low.
  - **Calculation Logic**: Sum the duration where `machine_working_condition = 'Stopped'` AND `current_level < low_level_limit`.
- **Sensor Needs Cleaning**: This is an anomaly detection pattern. A sensor likely needs cleaning if you observe the following pattern at least twice:
  - The material level remains constant (change = 0) for 5-10 consecutive minutes.
  - This period of no change is immediately followed by a sudden, large drop in the level reading (e.g., > 20mm in one minute).
  - **Level Change Formula**: `Level at (T-1) - Level at (T)`. Use window functions like `LAG()` to calculate this.

### OUTPUT REQUIREMENTS
You must generate a single JSON object containing two keys: "sql" and "chart". Do not add any explanatory text before or after the JSON object.

1.  **`sql` (string)**:
    - The query must be PostgreSQL-compliant.
    - Use `LOWER(column) LIKE LOWER('%value%')` for case-insensitive text matching.
    - Use `DATE_TRUNC('period', timestamp_column)` for time-based grouping (e.g., 'day', 'month').
    - Alias all derived columns (e.g., `AS total_revenue`).
    - Use window functions like `SUM(...) OVER (...)` or `LAG(...)` where appropriate, especially for complex calculations like "sensor needs cleaning".

2.  **`chart` (object)**:
    - **`type`**: Must be one of: `bar`, `line`, `pie`, `scatter`, `area`. Choose the most appropriate type.
        - `line` for trends over time.
        - `bar` for comparisons between categories.
        - `pie` for part-to-whole distributions.
    - **`x_axis`**: The column name for the x-axis.
    - **`y_axis`**: A list of column names for the y-axis.
    - **`labels`**: An object with `x` and `y` keys for axis labels.
    - **`title`**: A clear, descriptive title for the chart.

### EXAMPLE
**NLQ**: "Show monthly revenue vs cost per product line."
**Output**:
```json
{{
  "sql": "SELECT DATE_TRUNC('month', sale_date) AS month, product_line, SUM(revenue) AS total_revenue, SUM(cost) AS total_cost FROM Sales GROUP BY month, product_line ORDER BY month;",
  "chart": {{
    "type": "line",
    "x_axis": "month",
    "y_axis": ["total_revenue", "total_cost"],
    "labels": {{"x": "Month", "y": "Amount"}},
    "title": "Monthly Revenue vs Cost by Product Line"
  }}
}}


USER'S QUERY
NLQ: "{question}"

YOUR RESPONSE (JSON only)
    

"""

# Initialize the ChatOpenAI client
client = ChatOpenAI(model="gpt-4.1", openai_api_key=openapi)

# Format the HUMAN_PROMPT_TEMPLATE with the schema and the user's query
prompt = HUMAN_PROMPT_TEMPLATE.format(schema=schema_str, question=nlq_query)                                          
messages = [
        SystemMessage(content="You are an AI that generates SQL queries which will accepted by PostgreSQL and chart details from business-related natural language queries."),
        HumanMessage(content=prompt)
    ]
    
response = client.invoke(messages)
    
# Optionally, validate that the response is valid JSON:
try:
    result = json.loads(response.content)
    print(result)
except json.JSONDecodeError as e:
    raise ValueError("The response is not a valid JSON object.") from e

return result

throws error
JSONDecodeError                           Traceback (most recent call last)

<ipython-input-28-3460407561> in <cell line: 0>()
     89 try:
---> 90     result = json.loads(response.content)
     91     print(result)

3 frames

JSONDecodeError: Expecting value: line 1 column 1 (char 0)


The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)

<ipython-input-28-3460407561> in <cell line: 0>()
     91     print(result)
     92 except json.JSONDecodeError as e:
---> 93     raise ValueError("The response is not a valid JSON object.") from e
     94 
     95 return result

ValueError: The response is not a valid JSON object.