import asyncio
import json
import re
from typing import Dict, Any
from pydantic import BaseModel, Field

from community.langchain_community.utilities.tableau.utils import http_post

class AnalyzeDataInputs(BaseModel):
    """Describes model inputs for usage of the analyze_data tool"""
    query: str = Field(..., description="Take the user query and format it as a SQL query. Do not insert unnecessary COUNT or COUNTDISTINCT unless asked to do so.")

    class Config:
        """Configuration options for the Pydantic model."""
        json_schema_extra = {
            "example": {
                "query": "show me average discount, total sales and profits by region sorted by profit"
            }
        }

async def query_vds(api_key: str, datasource_luid: str, url: str, query: Dict[str, Any]) -> Dict[str, Any]:
    """
    Queries Tableau's VizQL Data Service with a parameterized query in the payload.

    Args:
        api_key (str): The API key for authentication.
        datasource_luid (str): The LUID of the data source to query.
        url (str): The base URL for the Tableau server.
        query (Dict[str, Any]): The query parameters for the data request.

    Returns:
        Dict[str, Any]: The queried data sets as a dictionary if successful.

    Raises:
        RuntimeError: If the request fails or returns an error status code.
    """
    full_url = f"{url}/api/v1/vizql-data-service/query-datasource"

    payload = {
        "datasource": {
            "datasourceLuid": datasource_luid
        },
        "query": query
    }

    headers = {
        'X-Tableau-Auth': api_key,
        'Content-Type': 'application/json'
    }

    response = await http_post(endpoint=full_url, headers=headers, payload=payload)

    # Check if the request was successful (status code 200)
    if response['status'] == 200:
        return response['data']
    else:
        error_message = (
            f"Failed to query data source via Tableau VizQL Data Service. "
            f"Status code: {response['status']}. Response: {response['data']}"
        )
        raise RuntimeError(error_message)

# sends requests to Tableau's VizQL Data Service with payload written by Agent
async def get_headlessbi_data(payload: str, url: str, api_key: str, datasource_luid: str):
    """
    Returns a dictionary containing a data consisting of formatted markdown and a query plan
    describing the reasoning and steps the model performed in order to generate the query payload
    """
    agent_response = get_payload(payload)
    vds_payload = agent_response["payload"]
    query_plan = agent_response["query_plan"]

    # when data is available return it and the reasoning behind the query
    if vds_payload:
        headlessbi_data = await query_vds(
            api_key = api_key,
            datasource_luid = datasource_luid,
            url = url,
            query = vds_payload
        )

        # Convert to JSON string
        markdown_table = json_to_markdown(headlessbi_data['data'])

        # response includes markdown table with data + query plan
        return {
            "query_plan": query_plan,
            "data": markdown_table
        }
    # when the LLM cannot generate a query, the reasoning will explain why
    else:
        return {
            "query_plan": query_plan,
            "data": None
        }

# separates the written payload from the Agent's reasoning
def get_payload(output):
    """
    Extracts the LLM generated payload to query Tableau VizQL Data Service on behalf of
    the end user. It also separates the behavioral reasoning generated by the model, inserting
    both generations into individual dictionary keys
    """

    # LLM reasoning
    query_plan = output.split('JSON_payload')[0]

    # parse LLM output and query headless BI
    parsed_output = output.split('JSON_payload')[1]

    match = re.search(r'{.*}', parsed_output, re.DOTALL)
    if match:
        json_string = match.group(0)
        payload = json.loads(json_string)

        return {
            "query_plan": query_plan,
            "payload": payload
        }
    else:
        # for when no payload is possible to generate
        return {
            "query_plan": query_plan
        }

# convert JSON responses to formatted markdown
def json_to_markdown(json_data):
    """
    Parses a JSON response from Tableau's VizQL Data Service into formatted markdown
    """
    # Parse the JSON data if it's a string
    if isinstance(json_data, str):
        json_data = json.loads(json_data)
    # Check if the JSON data is a list and not empty
    if not isinstance(json_data, list) or not json_data:
        raise ValueError(f"Invalid JSON data, you may have an error or if the array is empty then it was not possible to resolve the query your wrote: {json_data}")

    # Extract headers from the first dictionary
    headers = json_data[0].keys()

    # Create the Markdown table header
    markdown_table = "| " + " | ".join(headers) + " |\n"
    markdown_table += "| " + " | ".join(['---'] * len(headers)) + " |\n"

    # Add each row to the Markdown table
    for entry in json_data:
        row = "| " + " | ".join(str(entry[header]) for header in headers) + " |"
        markdown_table += row + "\n"

    return markdown_table

# request metadata of declared datasource
async def query_vds_metadata(api_key: str, datasource_luid: str, url: str) -> Dict[str, Any]:
    """
    Gets metadata from the VizQL Data Service endpoint for the specified data source.

    Args:
        api_key (str): The API key for authentication.
        datasource_luid (str): The LUID of the data source to query.
        url (str): The base URL for the Tableau server.

    Returns:
        Dict[str, Any]: The metadata of the queried data source.

    Raises:
        RuntimeError: If the request fails or returns an error status code.
    """
    full_url = f"{url}/api/v1/vizql-data-service/read-metadata"

    payload = {
        "datasource": {
            "datasourceLuid": datasource_luid
        }
    }

    headers = {
        'X-Tableau-Auth': api_key,
        'Content-Type': 'application/json'
    }

    response = await http_post(endpoint=full_url, headers=headers, payload=payload)

    # Check if the request was successful (status code 200)
    if response['status'] == 200:
        return response['data']
    else:
        error_message = (
            f"Failed to obtain data source metadata from VizQL Data Service. "
            f"Status code: {response['status']}. Response: {response['data']}"
        )
        raise RuntimeError(error_message)

# extract column or field values
async def get_values(api_key: str, url: str, datasource_luid: str, caption: str):
    """
    Returns the available members or column values of a data source field
    """

    column_values = {'fields': [{'fieldCaption': caption}]}
    output = await query_vds(
        api_key = api_key,
        datasource_luid = datasource_luid,
        url = url,
        query = column_values
    )
    if output is None:
        return None
    sample_values = [list(item.values())[0] for item in output['data']][:4]
    return sample_values

# obtains datasource metadata to augment the tool prompt
async def augment_datasource_metadata(api_key: str, url: str, datasource_luid: str, prompt: Dict[str, str]):
    """
    Enhances the provided prompt (expecting a key called "data_model") with metadata
    describing a Tableau data source such that an Agent can correctly write queries to meet
    the user's needs with regards to fields, filters and calculations
    """

    # get metadata from VizQL Data Service endpoint
    datasource_metadata = await query_vds_metadata(
        api_key=api_key,
        url=url,
        datasource_luid=datasource_luid
    )

    # Prepare a list of coroutines for fields that are of type 'STRING'
    string_fields_coroutines = []
    for field in datasource_metadata['data']:
        # Delete the specified keys
        del field['fieldName']
        del field['logicalTableId']

        # If the field is of type 'STRING', prepare the coroutine for getting values
        if field['dataType'] == 'STRING':
            coroutine = get_values(
                api_key=api_key,
                url=url,
                datasource_luid=datasource_luid,
                caption=field['fieldCaption']
            )
            string_fields_coroutines.append(coroutine)

    # Run all coroutines concurrently and wait for all to finish
    string_values_results = await asyncio.gather(*string_fields_coroutines)

    # Update the fields with the sample values
    for field, string_values in zip([field for field in datasource_metadata['data'] if field['dataType'] == 'STRING'], string_values_results):
        field['sampleValues'] = string_values

    # add the datasource metadata of the connected datasource to the system prompt
    prompt['data_model'] = datasource_metadata

    return json.dumps(prompt)
