import os
import duckdb
import groq
import sqlparse
import json
from typing import Dict, Any, List
import pandas as pd
from pathlib import Path
import gradio as gr
import shutil
from dotenv import load_dotenv
load_dotenv()
userdata = os.environ


def chat_with_groq(client:groq.Groq,
                   prompt:str,
                   model:str,
                   response_format: Dict[str, str]
                   )-> Any:
                   """
                   The function here works for the prompting the model for the qroq based model with natural language to
                   geegnrate the appropriate SQL query for the provided data and prompt. So with this output, can then be used to create proper
                   SQL queries and aggregation

                   Args:
                    client: This represents the Groq client which will be used top perform the conversion of the natural language to SQL.
                    prompt: Bases on the user message provided, the prompt will then be converted to SQL for future use for the data querying.
                    model: The Model here represents the based open source model that is used to make the query on. examples include llama3-70b-8192 or llama3-8b-8192 find out more here https://console.groq.com/docs/models
                    response_format: As model has the possibility of supporting JSON and MD outputs, this is a dictionary which specifies the the required return response format. by default this for the model is usually targetted
                    to returning an MD data but in the case of NL2SQL, then we coule work with SQL

                   Returns:
                   the function returns some data based on the response_format response form the model's response
                   """
                   completion = client.chat.completions.create(
                       model = model,
                       temperature = 0,
                       messages = [
                           {
                               "role":"user",
                               "content":prompt
                           }
                       ],
                       response_format=response_format
                   )
                  #  logger.info(f"Completion: {completion}")
                   return completion.choices[0].message.content

def execute_duckdb_query(query:str)->pd.DataFrame:
    """
    Execute a DuckDB query and return the result as a pandas DataFrame.

    Args:
        query (str): The DuckDB query to execute.

    Returns:
        pd.DataFrame: The result of the query as a pandas DataFrame.
    """
    original_cwd = os.getcwd()
    print(f"PATH:{original_cwd}")
    os.chdir('data')
    print(f"PATH:{os.getcwd()}")

    try:
        conn = duckdb.connect(database=":memory:", read_only=False)
        query_result = conn.execute(query).fetch_df().reset_index()
        os.chdir(original_cwd)
        return query_result
    except Exception as e:
        print(f"Error: {e}")
        os.chdir(original_cwd)
        raise e
def get_summarization(client:groq.Groq,
                      use_question:str,
                      df:pd.DataFrame,
                      model:str)->Any:
                      """
                      For this query, the user input is better summarized around the provided Dataframe. This give a better contextual structure for the user to perfom the action

                      Args:
                        client: This represents the Groq client which will be used top perform the conversion of the natural language to SQL.
                        use_question: Bases on the user message provided, the prompt will then be converted to SQL for future use for the data querying.
                        model: The Model here represents the based open source model that is used to make the query on. examples include llama3-70b-8192 or llama3-8b-8192 find out more here https://console.groq.com/docs/models
                        df: this is a pandas dataframe which containe the database daat which will then be used to summarize the base query around it, making the prompt more realistic.

                      Returns:
                      the function returns some json response form the model's response
                      """
                      prompt= '''
                      A user asked the following question pertaining to local database tables:

                      {user_question}

                      To answer the question, a dataframe was returned:

                      Dataframe:
                      {df}

                      In a few sentences, summarize the data in the table as it pertains to the original user question. Avoid qualifiers like "based on the data" and do not comment on the structure or metadata of the table itself
                      '''.format(user_question = use_question, df = df)
                      return chat_with_groq(client,prompt,model,None)


import re
from datetime import datetime
import numpy as np

def identify_value_datatype_to_SQL_DEF(value) -> str:

    pattern_yyyy_mm_dd = r'^\d{4}-\d{2}-\d{2}$'
    pattern_mm_dd_yyyy = r'^\d{2}/\d{2}/\d{4}$'
    pattern_dd_mm_yyyy = r'^\d{2}-\d{2}-\d{4}$'

    if pd.isna(value):
        return 'VARCHAR(255)'

    if isinstance(value, (int, np.int64)):
        return 'INTEGER'


    elif isinstance(value, (float, np.float64)):
        return 'FLOAT'


    elif isinstance(value, bool):
        return 'BOOLEAN'


    elif isinstance(value, (pd.Timestamp, datetime)):
        return 'DATETIME'


    elif isinstance(value, str):
        if re.match(pattern_yyyy_mm_dd, value) or re.match(pattern_mm_dd_yyyy, value) or re.match(pattern_dd_mm_yyyy, value):
            return 'DATE'


        try:
            pd.to_datetime(value)
            return 'DATETIME'
        except ValueError:
            pass


    elif isinstance(value, list):
        if all(isinstance(item, (int, np.int64)) for item in value):
            return 'ARRAY<INTEGER>'
        elif all(isinstance(item, str) for item in value):
            return 'ARRAY<VARCHAR(255)>'


    elif isinstance(value, dict):
        return 'JSON'

    elif isinstance(value, (str, bytes)):
        if value.startswith('POINT'):
            return 'GEOMETRY(POINT)'
        elif value.startswith('LINESTRING'):
            return 'GEOMETRY(LINESTRING)'
        elif value.startswith('POLYGON'):
            return 'GEOMETRY(POLYGON)'

    return 'VARCHAR(255)'

def identify_column_datatypes_to_SQL_DEF(df: pd.DataFrame,
                                         api_key:str,
                                         model:str
                                         ) -> dict:
    column_types = []

    for column in df.columns:
        # Get non-null values
        non_null_values = df[column].dropna()

        if len(non_null_values) > 0:
            # Use the first non-null value to determine the type
            sample_value = non_null_values.iloc[0]
            column_types.append(f"{column} - {identify_value_datatype_to_SQL_DEF(sample_value)}")
        else:
            # If all values are null, default to VARCHAR
            column_types.append(f"{column} - VARCHAR(255)")

    column_datatypes= "\n".join(element for element in column_types)
    client = groq.Groq(api_key=api_key,)
    full_prompt = """
    You are a Database Query Advisor. Your task is to provide descriptions for given columns and their datatypes.

    Input Data:
    {columns}

    Instructions:
    1. Respond with a valid JSON document.
    2. For each column, provide a description and its datatype in this format:
      {{"column_name": {{"description": "<description here>", "dtype": "<datatype here>"}}}}
    3. If the data provided is invalid or insufficient, respond with:
      {{"error": "<explanation of the issue>"}}
    4. Ensure each column name is directly associated with its description and datatype.
    5. Return the entire output on a single line without line breaks.
    6. Keep your response simple and straightforward; avoid complex structures.

    Example Output:
    {{"column1": {{"description": "Unique identifier for each record", "dtype": "INTEGER"}}, "column2": {{"description": "Name of the product", "dtype": "VARCHAR(255)"}}}}

    Remember: Your entire response must be valid JSON and on a single line.
    """.format(columns=column_datatypes)

    response = chat_with_groq(client,full_prompt,model,{
          "type":"json_object"
      })
    print(f"OUTPUT:{response}")
    response = json.loads(response)

    formatted_data = {
        f"  {column} - {response[column]['dtype']}": response[column]['description']
        for column in response
    }
    return " \n".join(f"{key}: {value}" for key, value in formatted_data.items())

def join_with_and(items):
    if not items:
        return ""
    if len(items) == 1:
        return items[0]
    return ', '.join(items[:-1]) + ' and ' + items[-1]


import os
from pathlib import Path
from typing import List

base_prompt =  """
    You are Groq Advisor, and you are tasked with generating SQL queries for DuckDB based on user questions about data stored in two tables derived from CSV files:

    {table_description}

    Given a user's question about this data, write a valid DuckDB SQL query that accurately extracts or calculates the requested information from these tables and adheres to SQL best practices for DuckDB, optimizing for readability and performance where applicable.
    Make sure that you do not query a tables that does not exist. Ensure that a query provided comes from the tables provided.
    When asked about quantitative values, ensure that you look at the numerical data types and not the string data types such as when asked about top purchase, the amount can be a field to be looked at.
    When queries about id, ensure that the relevant field in the data which represents an id is looked at and queried.

    Here are some tips for writing DuckDB queries:
    * DuckDB syntax requires querying from the .csv file itself, i.e. {tables}. For example: SELECT * FROM sample.csv as sample
    * Because we are using .csv file always make sure the .csv is attached to the file being queried
    * All tables referenced MUST be aliased
    * DuckDB does not implicitly include a GROUP BY clause
    * CURRENT_DATE gets today's date
    * Aggregated fields like COUNT(*) must be appropriately named


    Question:
    --------
    {user_question}
    --------
    Reminder: Generate a DuckDB SQL to answer to the question:
    * respond as a valid JSON Document
    * [Best] If the question can be answered with the available tables: {{"sql": <sql here>}}
    * If the question cannot be answered with the available tables: {{"error": <explanation here>}}
    * Ensure that the entire output is returned on only one single line
    * Keep your query as simple and straightforward as possible; do not use subqueries
    """
table_description = """"""
tables_string = """"""
table_1 = """"""
table_1_wt_xt = """"""
user_question = """"""

# And some rules for querying the dataset:
# * Never include employee_id in the output - show employee name instead

# Also note that:
# * Valid values for product_name include 'Tesla','iPhone' and 'Humane pin'


def upload_file(files) -> List[str]:
    # will have to change to the private system is initiializes
    model = "llama3-8b-8192"
    api_key:str=userdata.get("GROQ_API_KEY")
    data_dir = Path("data")
    data_dir.mkdir(parents=True, exist_ok=True)
    if type(files) == str:
        files = [files]
    stored_paths = []
    stored_table_descriptions = []
    tables = []
    for file in files:
        filename = Path(file.name).name
        path = data_dir / filename

        # Copy the content of the temporary file to our destination
        with open(file.name, "rb") as source, open(path, "wb") as destination:
            destination.write(source.read())

        stored_paths.append(str(path.absolute()))
        table_description = identify_column_datatypes_to_SQL_DEF(pd.read_csv(path),api_key,model)
        desc = "Table: " + filename + "\n Columns:\n" + table_description
        stored_table_descriptions.append(desc)
        tables.append(filename)
    # constructing a string
    tables_string = join_with_and(tables)
    final = "\n".join(stored_table_descriptions)
    table_1_wt_xt = tables[0].split('.')[0]
    table_description = final
    tables_string = tables_string
    table_1 = tables[0]
    table_1_wt_xt = table_1_wt_xt
    return final

def user_prompt_sanitization(user_prompt:str)->str:
  guide = """
            You are an AI assistant specializing in database queries. Your task is to interpret user questions about a database and refine them based on the available table structures. Use the following table descriptions to guide your responses:

            Database Schema:
            {database_schema}
            
            Instructions:
            1. Interpret the user's question and identify the relevant tables and columns.
            2. Refine the query to use correct table and column names as per the schema.
            3. Ensure all IDs are properly referenced (e.g., user_id instead of just id).
            4. For quantitative queries (e.g., "top 5"), specify the ordering criteria.
            5. Infer necessary joins between tables when the query spans multiple tables.
            6. Provide a clear, concise refinement of the user's query that accurately reflects the database structure.
            
            Now, please refine the following user query:
            {user_question}
            
         """

  formatted_guide = guide.format(database_schema=table_description,user_question=user_prompt)
  api_key:str=userdata.get("GROQ_API_KEY")
  client = groq.Groq(api_key=api_key)
  return chat_with_groq(client,formatted_guide,"llama3-70b-8192",None)

def queryModel(user_prompt:str,model:str = "llama3-70b-8192",api_key:str=userdata.get("GROQ_API_KEY")):
  client = groq.Groq(api_key=api_key)
  user_prompt = user_prompt_sanitization(user_prompt)
  print(user_prompt)
  full_prompt = base_prompt.format(
      user_question=user_prompt,
      table_description=table_description,
      tables=tables_string,
      table_1=table_1,
      table_1_wt_xt=table_1_wt_xt
  )
  try:
     response = chat_with_groq(client,full_prompt,model,{
      "type":"json_object"
     })
  except Exception as e:
    return [(
        "Groq Advisor",
        "Error: " + str(e)
    )]
  response = json.loads(response)
  if "sql" in response:
    sql_query = response["sql"]
    try:
      results_df = execute_duckdb_query(sql_query)
    except Exception as e:
      return [(
          "Groq Advisor",
          "Error: " + str(e)
      )]

    fotmatted_sql_query = sqlparse.format(sql_query, reindent=True, keyword_case='upper')
    # print(f"SQL Query: {fotmatted_sql_query}")
    # print(results_df.to_markdown())
    query_n_results = "SQL Query: " + fotmatted_sql_query + "\n\n" + results_df.to_markdown()
    summarization = get_summarization(client,user_prompt,results_df,model)
    query_n_results += "\n\n" + summarization

    return [(
        "Groq Advisor",
        query_n_results
    )]
  elif "error" in response:
    return [(
        "Groq Advisor",
        "Error: " + response["error"]
    )]
  else:
    return [(
        "Groq Advisor",
        "Error: Unknown error"
    )]

with gr.Blocks() as demo:
    gr.Markdown("# CSV Database Query Interface")

    with gr.Tab("Upload CSV"):
        file_output = gr.File(file_count="multiple", label="Upload your CSV files")
        upload_button = gr.Button("Load CSV Files")
        upload_output = gr.Textbox(label="Upload Status", lines=5)

        upload_button.click(upload_file, inputs=file_output, outputs=upload_output)
    with gr.Tab("Query Interface"):
        chatbot = gr.Chatbot()
        with gr.Row():
            user_input = gr.Textbox(label="Enter your question")
            submit_button = gr.Button("Submit")
        submit_button.click(queryModel, inputs=[user_input], outputs=chatbot)



demo.launch(
    share=True
)