from typing import Any, Optional
from langchain_core.language_models.llms import LLM
from datetime import datetime
import concurrent.futures
import json

from .general import parse_list
from .timbr_utils import get_datasources, get_tags, get_concepts, get_concept_properties, validate_sql, get_properties_description, get_relationships_description, cache_with_version_check, encrypt_prompt, get_ontology_description
from .prompt_service import (
    get_determine_concept_prompt_template,
    get_generate_sql_prompt_template,
    get_generate_sql_reasoning_prompt_template,
    get_qa_prompt_template
)
from .. import config

def _clean_snowflake_prompt(prompt: Any) -> None:
    import re

    def clean_func(prompt_content: str) -> str:
        raw = prompt_content
        # 1. Normalize Windows/Mac line endings → '\n'
        raw = raw.replace('\r\n', '\n').replace('\r', '\n')

        # 2. Collapse any multiple blank lines → single '\n'
        raw = re.sub(r'\n{2,}', '\n', raw)

        # 3. Convert ALL real '\n' → literal backslash-n
        raw = raw.replace('\n', '\\n')

        # 4. Normalize curly quotes to straight ASCII
        raw = (raw
            .replace('’', "'")
            .replace('‘', "'")
            .replace('“', '"')
            .replace('”', '"'))

        # 5. Collapse any accidental double-backticks → single backtick
        raw = raw.replace('``', '`')

        # 6. Escape ALL backslashes so '\\n' survives as two chars
        raw = raw.replace('\\', '\\\\')

        # 7. Escape single-quotes for SQL string literal
        raw = raw.replace("'", "''")

        # 8. Escape double-quotes for SQL string literal
        raw = raw.replace('"', '\\"')

        return raw

    prompt[0].content = clean_func(prompt[0].content)  # System message
    prompt[1].content = clean_func(prompt[1].content)  # User message


def _call_llm_with_timeout(llm: LLM, prompt: Any, timeout: int = 120) -> Any:
    """
    Call LLM with timeout to prevent hanging.
    
    Args:
        llm: The LLM instance
        prompt: The prompt to send
        timeout: Timeout in seconds (default: 120)
        
    Returns:
        LLM response
        
    Raises:
        TimeoutError: If the call takes longer than timeout seconds
        Exception: Any other exception from the LLM call
    """
    def _llm_call():
        return llm.invoke(prompt)
    
    with concurrent.futures.ThreadPoolExecutor() as executor:
        future = executor.submit(_llm_call)
        try:
            return future.result(timeout=timeout)
        except concurrent.futures.TimeoutError:
            raise TimeoutError(f"LLM call timed out after {timeout} seconds")
        except Exception as e:
            raise e

MEASURES_DESCRIPTION = "The following columns are calculated measures and can only be aggregated with an aggregate function: COUNT/SUM/AVG/MIN/MAX (count distinct is not allowed)"
TRANSITIVE_RELATIONSHIP_DESCRIPTION = "Transitive relationship columns allow you to access data through multiple relationship hops. These columns follow the pattern `<relationship_name>[<table_name>*<number>].<column_name>` where the number after the asterisk (*) indicates how many relationship levels to traverse. For example, `acquired_by[company*4].company_name` means 'go through up to 4 levels of the acquired_by relationship to get the company name', while columns ending with '_transitivity_level' indicate the actual relationship depth (Cannot be null or 0 - level 1 represents direct relationships, while levels 2, 3, 4, etc. represent indirect relationships through multiple hops. To filter by relationship type, use `_transitivity_level = 1` for direct relationships only, `_transitivity_level > 1` for indirect relationships only."


def _prompt_to_string(prompt: Any) -> str:
    prompt_text = ''
    if isinstance(prompt, str):
        prompt_text = prompt
    elif isinstance(prompt, list):
        for message in prompt:
            if hasattr(message, "content"):
                prompt_text += message.type + ": " + message.content + "\n"
            else:
                prompt_text += str(message)
    else:
        prompt_text = str(prompt)
    return prompt_text.strip()


def _calculate_token_count(llm: LLM, prompt: str | list[Any]) -> int:
    """
    Calculate the token count for a given prompt text using the specified LLM.
    Falls back to basic if the LLM doesn't support token counting.
    """
    import tiktoken
    token_count = 0

    encoding = None
    try:
        if hasattr(llm, 'client') and hasattr(llm.client, 'model_name'):
            encoding = tiktoken.encoding_for_model(llm.client.model_name)
    except Exception as e:
        print(f"Error with primary token counting: {e}")
        pass

    try:
        if encoding is None:
            encoding = tiktoken.get_encoding("cl100k_base")
        if isinstance(prompt, str):
            token_count = len(encoding.encode(prompt))
        else:
            prompt_text = _prompt_to_string(prompt)
            token_count = len(encoding.encode(prompt_text))
    except Exception as e2:
        #print(f"Error calculating token count with fallback method: {e2}")
        pass

    return token_count
    

def _get_response_text(response: Any) -> str:
    if hasattr(response, "content"):
        response_text = response.content

        # Handle Databricks gpt-oss type of responses (having list of dicts with type + summary for reasoning or type + text for result)
        if isinstance(response_text, list):
            response_text = next(filter(lambda x: x.get('type') == 'text', response.content), None)
        if isinstance(response_text, dict):
            response_text = response_text.get('text', '')
    elif isinstance(response, str):
        response_text = response
    else:
        raise ValueError("Unexpected response format from LLM.")

    if "QUESTION VALIDATION ERROR:" in response_text:
        err = response_text.split("QUESTION VALIDATION ERROR:", 1)[1].strip()
        raise ValueError(err)

    return response_text


def _extract_usage_metadata(response: Any) -> dict:
    """
    Extract usage metadata from LLM response across different providers.
    
    Different providers return usage data in different formats:
    - OpenAI/AzureOpenAI: response.response_metadata['token_usage'] or response.usage_metadata
    - Anthropic: response.response_metadata['usage'] or response.usage_metadata
    - Google/VertexAI: response.usage_metadata
    - Bedrock: response.response_metadata['usage'] or response.response_metadata (direct ResponseMetadata)
    - Snowflake: response.response_metadata['usage']
    - Databricks: response.usage_metadata or response.response_metadata
    """
    usage_metadata = {}
    
    # Try to get response_metadata first (most common)
    if hasattr(response, 'response_metadata') and response.response_metadata:
        resp_meta = response.response_metadata
        
        # Check for 'usage' key (Anthropic, Bedrock, Snowflake)
        if 'usage' in resp_meta:
            usage_metadata = resp_meta['usage']
        # Check for 'token_usage' key (OpenAI/AzureOpenAI)
        elif 'token_usage' in resp_meta:
            usage_metadata = resp_meta['token_usage']
        # Check for direct token fields in response_metadata (some Bedrock responses)
        elif any(key in resp_meta for key in ['input_tokens', 'output_tokens', 'total_tokens', 
                                                'prompt_tokens', 'completion_tokens']):
            usage_metadata = {
                k: v for k, v in resp_meta.items() 
                if k in ['input_tokens', 'output_tokens', 'total_tokens', 
                        'prompt_tokens', 'completion_tokens']
            }
    
    # Try usage_metadata attribute (Google, VertexAI, some others)
    if not usage_metadata and hasattr(response, 'usage_metadata') and response.usage_metadata:
        usage_meta = response.usage_metadata
        if isinstance(usage_meta, dict):
            # If it has a nested 'usage' key
            if 'usage' in usage_meta:
                usage_metadata = usage_meta['usage']
            else:
                usage_metadata = usage_meta
        else:
            # Handle case where usage_metadata is an object with attributes
            usage_metadata = {
                k: getattr(usage_meta, k) 
                for k in dir(usage_meta) 
                if not k.startswith('_') and not callable(getattr(usage_meta, k))
            }
    
    # Try direct usage attribute (fallback)
    if not usage_metadata and hasattr(response, 'usage') and response.usage:
        usage = response.usage
        if isinstance(usage, dict):
            if 'usage' in usage:
                usage_metadata = usage['usage']
            else:
                usage_metadata = usage
        else:
            # Handle case where usage is an object with attributes
            usage_metadata = {
                k: getattr(usage, k) 
                for k in dir(usage) 
                if not k.startswith('_') and not callable(getattr(usage, k))
            }
    
    # Normalize token field names to standard format
    # Different providers use different names: input_tokens vs prompt_tokens, etc.
    if usage_metadata:
        normalized = {}
        
        # Map various input token field names
        if 'input_tokens' in usage_metadata:
            normalized['input_tokens'] = usage_metadata['input_tokens']
        elif 'prompt_tokens' in usage_metadata:
            normalized['input_tokens'] = usage_metadata['prompt_tokens']
        
        # Map various output token field names
        if 'output_tokens' in usage_metadata:
            normalized['output_tokens'] = usage_metadata['output_tokens']
        elif 'completion_tokens' in usage_metadata:
            normalized['output_tokens'] = usage_metadata['completion_tokens']
        
        # Map total tokens
        if 'total_tokens' in usage_metadata:
            normalized['total_tokens'] = usage_metadata['total_tokens']
        elif 'input_tokens' in normalized and 'output_tokens' in normalized:
            # Calculate total if not provided
            normalized['total_tokens'] = normalized['input_tokens'] + normalized['output_tokens']
        
        # Keep any other metadata fields that don't conflict
        for key, value in usage_metadata.items():
            if key not in ['input_tokens', 'prompt_tokens', 'output_tokens', 
                          'completion_tokens', 'total_tokens']:
                normalized[key] = value
        
        return normalized if normalized else usage_metadata
    
    return usage_metadata

def filter_list_by_ontology(concepts_list, ontology) -> list:
    ontology_specific_concepts = concepts_list
    
    if concepts_list is not None and isinstance(concepts_list, list):
        ontology_specific_concepts = []
        for concept in concepts_list:
            if "." not in concept:
                ontology_specific_concepts.append(concept)
            elif concept.startswith(f"{ontology}."):
                ontology_specific_concepts.append(concept.split(".", 1)[1])
    
    return ontology_specific_concepts

def determine_concept(
    question: str,
    llm: LLM,
    conn_params: dict,
    concepts_list: Optional[list] = None,
    views_list: Optional[list] = None,
    include_logic_concepts: Optional[bool] = False,
    include_tags: Optional[str] = None,
    should_validate: Optional[bool] = False,
    retries: Optional[int] = 3,
    note: Optional[str] = '',
    debug: Optional[bool] = False,
    timeout: Optional[int] = None,
) -> dict[str, Any]:
    usage_metadata = {}
    determined_concept_name = None
    identify_concept_reason = None
    schema = 'dtimbr'
    
    # Use config default timeout if none provided
    if timeout is None:
        timeout = config.llm_timeout
    
    determine_concept_prompt = get_determine_concept_prompt_template(conn_params)

    ontologies_conn_params = {}
    #ontologies = [o.strip().lower() for o in conn_params.get("ontology").split(",")]
    ontologies = parse_list(conn_params.get("ontology"))

    for ontology in ontologies:
        ontology_conn_param =  conn_params.copy()
        ontology_conn_param["ontology"] = ontology
        ontologies_conn_params[ontology] = ontology_conn_param

    concepts_desc_arr = []
    ontologies_concepts_and_views = {}
    candidates = []

    for ontology in ontologies_conn_params.keys():
        tags = get_tags(conn_params=ontologies_conn_params[ontology], include_tags=include_tags)
        ontology_description, domain_description = get_ontology_description(ontologies_conn_params[ontology])

        ontology_specific_concepts = filter_list_by_ontology(concepts_list, ontology)
        ontology_specific_views = filter_list_by_ontology(views_list, ontology)

        concepts_and_views = get_concepts(
            conn_params=ontologies_conn_params[ontology],
            concepts_list=ontology_specific_concepts,
            views_list=ontology_specific_views,
            include_logic_concepts=include_logic_concepts,
        )

        if concepts_and_views:
            ontologies_concepts_and_views[ontology] = concepts_and_views

            formatted_ontology_desc = f"-- Schema `{ontology}`" 
            
            if ontology_description != "":
                cleaned_ontology_desc = ontology_description.replace("\r\n", " ").replace("\n", " ")
                formatted_ontology_desc += f" description: {cleaned_ontology_desc}"

            if domain_description != "":
                cleaned_domain_desc = domain_description.replace("\r\n", " ").replace("\n", " ")
                formatted_ontology_desc += f". Related Domains description: {cleaned_domain_desc}"
            concepts_desc_arr.append(formatted_ontology_desc + "\n")

            for item in concepts_and_views.values():
                item_name = item.get('concept')
                item_desc = item.get('description')
                item_tags = tags.get('concept_tags').get(item_name) if item.get('is_view') == 'false' else tags.get('view_tags').get(item_name)

                if item_tags:
                    item_tags = str(item_tags).replace('{', '').replace('}', '').replace("'", '')

                clean_prefix = ""
                prefix = ""

                if len(ontologies_conn_params) > 1:
                    clean_prefix = f"{ontology}."
                    prefix = f"`{ontology}`."

                candidates.append(clean_prefix + item_name)
                concept_verbose = prefix + f"`{item_name}`"
                if item_desc:
                    concept_verbose += f" (description: {item_desc})"
                if item_tags:
                    concept_verbose += f" [tags: {item_tags}]"
                    concepts_and_views[item_name]['tags'] = f"- Annotations and constraints: {item_tags}\n"

                concepts_desc_arr.append(concept_verbose)
            
            concepts_desc_arr.append('\n')

    if len(ontologies_concepts_and_views) == 0:
        raise Exception("No relevant concepts found for the query.")

    if len(ontologies_concepts_and_views) == 1 and len(list(ontologies_concepts_and_views.values())[0]) == 1:
        # If only one concept is provided, return it directly
        determined_concept_name = list(list(ontologies_concepts_and_views.values())[0].keys())[0]
    else:
        # Use LLM to determine the concept based on the question
        iteration = 0
        error = ''
        while determined_concept_name is None and iteration < retries:
            iteration += 1
            err_txt = f"\nLast try got an error: {error}" if error else ""
            prompt = determine_concept_prompt.format_messages(
                question=question.strip(),
                concepts="\n".join(concepts_desc_arr),
                note=(note or '') + err_txt,
            )

            # temporary fix to old prompts 
            if len(prompt) == 2:
                prompt[1].content = prompt[1].content.replace("no quotes", "no backtick quotes")
            
            apx_token_count = _calculate_token_count(llm, prompt)
            if "snowflake" in llm._llm_type:
                _clean_snowflake_prompt(prompt)

            try:
                response = _call_llm_with_timeout(llm, prompt, timeout=timeout)
            except TimeoutError as e:
                error = f"LLM call timed out: {str(e)}"
                raise Exception(error)
            except Exception as e:
                error = f"LLM call failed: {str(e)}"
                continue
            usage_metadata['determine_concept'] = {
                "approximate": apx_token_count,
                **_extract_usage_metadata(response),
            }
            if debug:
                usage_metadata['determine_concept']["p_hash"] = encrypt_prompt(prompt)

            # Try to parse as JSON first (with 'result' and 'reason' keys)
            try:
                parsed_response = _parse_json_from_llm_response(response)
                if isinstance(parsed_response, dict) and 'result' in parsed_response:
                    candidate = parsed_response.get('result', '').strip().replace("`", "").replace('"', "").lower()
                    identify_concept_reason = parsed_response.get('reason', None)
                else:
                    # Fallback to plain text if JSON doesn't have expected structure
                    candidate = _get_response_text(response).strip().replace("`", "").replace('"', "").lower()
            except (json.JSONDecodeError, ValueError):
                # If not JSON, treat as plain text (backwards compatibility)
                candidate = _get_response_text(response).strip().replace("`", "").replace('"', "").lower()
            
            if candidate not in candidates:

                if len(ontologies_conn_params) > 1:
                    for existing in candidates:
                        if existing.endswith("." + candidate):
                            candidate = existing

                if candidate not in candidates:            
                    error = f"Concept '{candidate}' not found in the list of concepts."
                    continue
            
            determined_concept_name = candidate
            error = ''

        if determined_concept_name is None and error != '':
            raise Exception(f"Failed to determine concept: {error}")

    ontology = list(ontologies_concepts_and_views.keys())[0]
    concepts_and_views = list(ontologies_concepts_and_views.values())[0]
    
    if "." in determined_concept_name:
        ontology = determined_concept_name.split(".")[0]
        determined_concept_name = determined_concept_name.split(".")[1]
        concepts_and_views = ontologies_concepts_and_views.get(ontology)    

    schema = 'vtimbr' if concepts_and_views.get(determined_concept_name).get('is_view') == 'true' else 'dtimbr'

    return {
        "concept_metadata": concepts_and_views.get(determined_concept_name) if determined_concept_name else None,
        "concept": determined_concept_name,
        "identify_concept_reason": identify_concept_reason,
        "schema": schema,
        "usage_metadata": usage_metadata,
        "ontology": ontology,
        "conn_params": ontologies_conn_params.get(ontology)
    }


def _build_columns_str(
    columns: list[dict],
    columns_tags: Optional[dict] = {},
    exclude: Optional[list] = None,
) -> str:
    columns_desc_arr = []
    for col in columns:
        full_name = col.get('name') or col.get('col_name') # When rel column, it can be `relationship_name[column_name]`
        col_name = col.get('col_name', '')

        if col_name.startswith("measure."):
            col_name = col_name.replace("measure.", "")

        if exclude and (col_name in exclude or any(col_name.endswith('.' + exc) for exc in exclude)):
            continue

        col_tags = str(columns_tags.get(col_name)) if columns_tags.get(col_name) else None
        if col_tags:
            col_tags = col_tags.replace('{', '').replace('}', '').replace("'", '').replace(": ", " - ").replace(",", ". ").strip()
        
        description = col.get('description') or  col.get('comment', '')

        data_type = col.get('data_type', 'string').lower() or 'string'

        col_meta = []
        if data_type:
            col_meta.append(f"type: {data_type}")
        if description:
            col_meta.append(f"description: {description}")
        if col_tags:
            col_meta.append(f"annotations and constraints: {col_tags}")

        col_meta_str = ', '.join(col_meta) if col_meta else ''
        if col_meta_str:
            col_meta_str = f" ({col_meta_str})"

        columns_desc_arr.append(f"`{full_name}`{col_meta_str}")

    return ", ".join(columns_desc_arr) if columns_desc_arr else ''


def _build_rel_columns_str(relationships: list[dict], columns_tags: Optional[dict] = {}, exclude_properties: Optional[list] = None) -> str:
    if not relationships:
        return ''
    rel_str_arr = []
    for rel_name in relationships:
        rel = relationships[rel_name]
        rel_description = rel.get('description', '')
        rel_description = f" which described as \"{rel_description}\"" if rel_description else ""
        rel_columns = rel.get('columns', [])
        rel_measures = rel.get('measures', [])
        
        if rel_columns:
            joined_columns_str = _build_columns_str(rel_columns, columns_tags=columns_tags, exclude=exclude_properties)
            rel_str_arr.append(f"- The following columns are part of {rel_name} relationship{rel_description}, and must be used as is wrapped with quotes: {joined_columns_str}")
        if rel_measures:
            joined_measures_str = _build_columns_str(rel_measures, columns_tags=columns_tags, exclude=exclude_properties)
            rel_str_arr.append(f"- {MEASURES_DESCRIPTION}, are part of {rel_name} relationship{rel_description}: {joined_measures_str}")
    
    return '.\n'.join(rel_str_arr) if rel_str_arr else ''


def _parse_sql_and_reason_from_llm_response(response: Any) -> dict:
    """
    Parse SQL & reason from LLM response. Handles both plain SQL strings and JSON format with 'result' and 'reason' keys.
    
    Returns:
        dict with 'sql' and 'reason' keys (reason may be None if not provided)
    """
    # Try to parse as JSON first
    try:
        parsed_json = _parse_json_from_llm_response(response)
        
        # Extract SQL from 'result' key and reason from 'reason' key
        if isinstance(parsed_json, dict) and 'result' in parsed_json:
            sql = parsed_json.get('result', '')
            reason = parsed_json.get('reason', None)
            
            # Clean the SQL
            sql = (sql
                   .replace("```sql", "")
                   .replace("```", "")
                   .replace('SELECT \n', 'SELECT ')
                   .replace(';', '')
                   .strip())
            
            return {'sql': sql, 'reason': reason}
    except (json.JSONDecodeError, ValueError):
        # If not JSON, treat as plain SQL string (backwards compatibility)
        pass
    
    # Fallback to plain text parsing
    response_text = _get_response_text(response)
    sql = (response_text
           .replace("```sql", "")
           .replace("```", "")
           .replace('SELECT \n', 'SELECT ')
           .replace(';', '')
           .strip())
    
    return {'sql': sql, 'reason': None}


def _get_active_datasource(conn_params: dict) -> dict:
    datasources = get_datasources(conn_params, filter_active=True)
    return datasources[0] if datasources else None


def _parse_json_from_llm_response(response: Any) -> dict:
    """
    Parse JSON from LLM response. Handles markdown code blocks and extracts valid JSON.
    
    Args:
        response: LLM response object
        
    Returns:
        dict containing parsed JSON
        
    Raises:
        json.JSONDecodeError: If response cannot be parsed as JSON
        ValueError: If response format is unexpected
    """
    response_text = _get_response_text(response)
    
    # Remove markdown code block markers if present
    content = response_text.strip()
    if content.startswith("```json"):
        content = content[7:]  # Remove ```json
    elif content.startswith("```"):
        content = content[3:]  # Remove ```
    
    if content.endswith("```"):
        content = content[:-3]  # Remove closing ```
    
    content = content.strip()
    
    # Parse and return JSON
    return json.loads(content)


def _evaluate_sql_enable_reasoning(
    question: str,
    sql_query: str,
    llm: LLM,
    conn_params: dict,
    timeout: int,
) -> dict:
    """
    Evaluate if the generated SQL correctly answers the business question.
    
    Returns:
        dict with 'assessment' ('correct'|'partial'|'incorrect') and 'reasoning'
    """
    generate_sql_reasoning_template = get_generate_sql_reasoning_prompt_template(conn_params)
    prompt = generate_sql_reasoning_template.format_messages(
        question=question.strip(),
        sql_query=sql_query.strip(),
    )

    apx_token_count = _calculate_token_count(llm, prompt)
    if hasattr(llm, "_llm_type") and "snowflake" in llm._llm_type:
        _clean_snowflake_prompt(prompt)
    
    response = _call_llm_with_timeout(llm, prompt, timeout=timeout)
    
    # Parse JSON response
    evaluation = _parse_json_from_llm_response(response)
    
    return {
        "evaluation": evaluation,
        "apx_token_count": apx_token_count,
        "usage_metadata": _extract_usage_metadata(response),
    }


@cache_with_version_check
def _build_sql_generation_context(
    conn_params: dict,
    schema: str,
    concept: str,
    concept_metadata: dict,
    graph_depth: int,
    include_tags: Optional[str],
    exclude_properties: Optional[list],
    db_is_case_sensitive: bool,
    max_limit: int,
) -> dict:
    """
    Prepare the complete SQL generation context by gathering all necessary metadata.
    
    This includes:
    - Datasource information
    - Concept properties (columns, measures, relationships)
    - Property tags
    - Building column/measure/relationship descriptions
    - Assembling the final context dictionary
    
    Returns:
        dict containing all context needed for SQL generation prompts
    """
    datasource_type = _get_active_datasource(conn_params).get('target_type')

    properties_desc = get_properties_description(conn_params=conn_params)
    relationships_desc = get_relationships_description(conn_params=conn_params)
  
    concept_properties_metadata = get_concept_properties(
        schema=schema,
        concept_name=concept,
        conn_params=conn_params,
        properties_desc=properties_desc,
        relationships_desc=relationships_desc,
        graph_depth=graph_depth
    )
    columns = concept_properties_metadata.get('columns', [])
    measures = concept_properties_metadata.get('measures', [])
    relationships = concept_properties_metadata.get('relationships', {})
    tags = get_tags(conn_params=conn_params, include_tags=include_tags).get('property_tags')

    columns_str = _build_columns_str(columns, columns_tags=tags, exclude=exclude_properties)
    measures_str = _build_columns_str(measures, tags, exclude=exclude_properties)
    rel_prop_str = _build_rel_columns_str(relationships, columns_tags=tags, exclude_properties=exclude_properties)

    if rel_prop_str:
        measures_str += f"\n{rel_prop_str}"

    # Determine if relationships have transitive properties
    has_transitive_relationships = any(
        rel.get('is_transitive')
        for rel in relationships.values()
    ) if relationships else False
    
    concept_description = f"- Description: {concept_metadata.get('description')}\n" if concept_metadata and concept_metadata.get('description') else ""
    concept_tags = concept_metadata.get('tags') if concept_metadata and concept_metadata.get('tags') else ""
    
    cur_date = datetime.now().strftime("%Y-%m-%d")
    
    # Build context descriptions
    sensitivity_txt = "- Ensure value comparisons are case-insensitive, e.g., use LOWER(column) = 'value'.\n" if db_is_case_sensitive else ""
    measures_context = f"- {MEASURES_DESCRIPTION}: {measures_str}\n" if measures_str else ""
    transitive_context = f"- {TRANSITIVE_RELATIONSHIP_DESCRIPTION}\n" if has_transitive_relationships else ""
    
    return {
        'cur_date': cur_date,
        'datasource_type': datasource_type or 'standard sql',
        'schema': schema,
        'concept': concept,
        'concept_description': concept_description or "",
        'concept_tags': concept_tags or "",
        'columns_str': columns_str,
        'measures_context': measures_context,
        'transitive_context': transitive_context,
        'sensitivity_txt': sensitivity_txt,
        'max_limit': max_limit,
    }


def _generate_sql_with_llm(
    question: str,
    llm: LLM,
    generate_sql_prompt: Any,
    current_context: dict,
    note: str,
    timeout: int,
    debug: bool = False,
) -> dict:
    """
    Generate SQL using LLM based on the provided context and note.
    This function is used for both initial SQL generation and regeneration with feedback.
    
    Args:
        current_context: dict containing datasource_type, schema, concept, concept_description,
                        concept_tags, columns_str, measures_context, transitive_context,
                        sensitivity_txt, max_limit, cur_date
        note: Additional instructions/feedback to include in the prompt
    
    Returns:
        dict with 'sql', 'is_valid', 'error', 'apx_token_count', 'usage_metadata', 'p_hash' (if debug)
    """
    prompt = generate_sql_prompt.format_messages(
        current_date=current_context['cur_date'],
        datasource_type=current_context['datasource_type'],
        schema=current_context['schema'],
        concept=f"`{current_context['concept']}`",
        description=current_context['concept_description'],
        tags=current_context['concept_tags'],
        question=question,
        columns=current_context['columns_str'],
        measures_context=current_context['measures_context'],
        transitive_context=current_context['transitive_context'],
        sensitivity_context=current_context['sensitivity_txt'],
        max_limit=current_context['max_limit'],
        note=note,
    )

    apx_token_count = _calculate_token_count(llm, prompt)
    if hasattr(llm, "_llm_type") and "snowflake" in llm._llm_type:
        _clean_snowflake_prompt(prompt)
    
    response = _call_llm_with_timeout(llm, prompt, timeout=timeout)
    
    # Parse response which now includes both SQL and reason
    parsed_response = _parse_sql_and_reason_from_llm_response(response)
    
    result = {
        "sql": parsed_response['sql'],
        "generate_sql_reason": parsed_response['reason'],
        "apx_token_count": apx_token_count,
        "usage_metadata": _extract_usage_metadata(response),
        "is_valid": True,
        "error": None,
    }
    
    if debug:
        result["p_hash"] = encrypt_prompt(prompt)
    
    
    return result

def handle_generate_sql_reasoning(
    sql_query: str,
    question: str,
    llm: LLM,
    conn_params: dict,
    schema: str,
    concept: str,
    concept_metadata: dict,
    include_tags: bool,
    exclude_properties: list,
    db_is_case_sensitive: bool,
    max_limit: int,
    reasoning_steps: int,
    note: str,
    graph_depth: int,
    usage_metadata: dict,
    timeout: int,
    debug: bool,
    previous_token_count
) -> tuple[str, int, str]:
    generate_sql_prompt = get_generate_sql_prompt_template(conn_params)
    context_graph_depth = graph_depth
    reasoned_sql = sql_query
    reasoned_sql_reason = None
    for step in range(reasoning_steps):
        try:
            # Step 1: Evaluate the current SQL
            eval_result = _evaluate_sql_enable_reasoning(
                question=question,
                sql_query=reasoned_sql,
                llm=llm,
                conn_params=conn_params,
                timeout=timeout,
            )
            
            usage_metadata[f'sql_reasoning_step_{step + 1}'] = {
                "approximate": eval_result['apx_token_count'],
                **eval_result['usage_metadata'],
            }
            
            evaluation = eval_result['evaluation']
            reasoning_status = evaluation.get("assessment", "partial").lower()
            reasoned_sql_reason = evaluation.get("reasoning", "")
            
            if reasoning_status == "correct":
                break
            
            # Step 2: Regenerate SQL with feedback
            evaluation_note = note + f"\n\nThe previously generated SQL: `{reasoned_sql}` was assessed as '{evaluation.get('assessment')}' because: {reasoned_sql_reason or '*could not determine cause*'}. Please provide a corrected SQL query that better answers the question: '{question}'.\n\nCRITICAL: Return ONLY the SQL query without any explanation or comments."
            
            # Increase graph depth for 2nd+ reasoning attempts, up to max of 3
            max_context_graph_depth = 3
            context_graph_depth = min(max_context_graph_depth, graph_depth)

            if (step >= 1 and type(previous_token_count) == int and previous_token_count > 0 and previous_token_count < 20000):
                context_graph_depth = min(max_context_graph_depth, context_graph_depth + 1)

            regen_result = _generate_sql_with_llm(
                question=question,
                llm=llm,
                generate_sql_prompt=generate_sql_prompt,
                current_context=_build_sql_generation_context(
                    conn_params=conn_params,
                    schema=schema,
                    concept=concept,
                    concept_metadata=concept_metadata,
                    graph_depth=context_graph_depth,
                    include_tags=include_tags,
                    exclude_properties=exclude_properties,
                    db_is_case_sensitive=db_is_case_sensitive,
                    max_limit=max_limit),
                note=evaluation_note,
                timeout=timeout,
                debug=debug,
            )
            
            reasoned_sql = regen_result['sql']
            reasoned_sql_reason = regen_result['generate_sql_reason']
            error = regen_result['error']

            step_key = f'generate_sql_reasoning_step_{step + 1}'
            usage_metadata[step_key] = {
                "approximate": regen_result['apx_token_count'],
                **regen_result['usage_metadata'],
            }
            previous_token_count = regen_result['apx_token_count']

            if debug and 'p_hash' in regen_result:
                usage_metadata[step_key]['p_hash'] = regen_result['p_hash']

            if error:
                raise Exception(error)
            
        except TimeoutError as e:
            raise Exception(f"LLM call timed out: {str(e)}")
        except Exception as e:
            print(f"Warning: LLM reasoning failed: {e}")
            break
    
    return reasoned_sql, context_graph_depth, reasoned_sql_reason

def handle_validate_generate_sql(
    sql_query: str,
    question: str,
    llm: LLM,
    conn_params: dict,
    generate_sql_prompt: Any,
    schema: str,
    concept: str,
    concept_metadata: dict,
    include_tags: bool,
    exclude_properties: list,
    db_is_case_sensitive: bool,
    max_limit: int,
    graph_depth: int,
    retries: int,
    timeout: int,
    debug: bool,
    usage_metadata: dict,
) -> tuple[bool, str, str]:
    is_sql_valid, error, sql_query = validate_sql(sql_query, conn_params)
    validation_attempt = 0
  
    while validation_attempt < retries and not is_sql_valid:
        validation_attempt += 1
        validation_err_txt = f"\nThe generated SQL (`{sql_query}`) was invalid with error: {error}. Please generate a corrected query that achieves the intended result." if error and "snowflake" not in llm._llm_type else ""

        regen_result = _generate_sql_with_llm(
            question=question,
            llm=llm,
            generate_sql_prompt=generate_sql_prompt,
            current_context=_build_sql_generation_context(
                conn_params=conn_params,
                schema=schema,
                concept=concept,
                concept_metadata=concept_metadata,
                graph_depth=graph_depth,
                include_tags=include_tags,
                exclude_properties=exclude_properties,
                db_is_case_sensitive=db_is_case_sensitive,
                max_limit=max_limit),
            note=validation_err_txt,
            timeout=timeout,
            debug=debug,
        )
        
        regen_error = regen_result['error']
        sql_query = regen_result['sql']

        validation_key = f'generate_sql_validation_regen_{validation_attempt}'
        usage_metadata[validation_key] = {
            "approximate": regen_result['apx_token_count'],
            **regen_result['usage_metadata'],
        }
        if debug and 'p_hash' in regen_result:
            usage_metadata[validation_key]['p_hash'] = regen_result['p_hash']

        if regen_error:
            raise Exception(regen_error)
        
        is_sql_valid, error, sql_query = validate_sql(sql_query, conn_params)

    return is_sql_valid, error, sql_query

def generate_sql(
        question: str,
        llm: LLM,
        conn_params: dict,
        concept: str,
        schema: Optional[str] = None,
        concepts_list: Optional[list] = None,
        views_list: Optional[list] = None,
        include_logic_concepts: Optional[bool] = False,
        include_tags: Optional[str] = None,
        exclude_properties: Optional[list] = None,
        should_validate_sql: Optional[bool] = config.should_validate_sql,
        retries: Optional[int] = 3,
        max_limit: Optional[int] = config.llm_default_limit,
        note: Optional[str] = '',
        db_is_case_sensitive: Optional[bool] = False,
        graph_depth: Optional[int] = 1,
        enable_reasoning: Optional[bool] = False,
        reasoning_steps: Optional[int] = 2,
        debug: Optional[bool] = False,
        timeout: Optional[int] = None,
    ) -> dict[str, str]:
    usage_metadata = {}
    concept_metadata = None
    reasoning_status = 'correct'

    # Use config default timeout if none provided
    if timeout is None:
        timeout = config.llm_timeout
    
    if concept and concept != "" and (schema is None or schema != "vtimbr"):
        concepts_list = [concept]
    elif concept and concept != "" and schema == "vtimbr":
        views_list = [concept]

    determine_concept_res = determine_concept(
        question=question,
        llm=llm,
        conn_params=conn_params,
        concepts_list=concepts_list,
        views_list=views_list,
        include_logic_concepts=include_logic_concepts,
        include_tags=include_tags,
        should_validate=should_validate_sql,
        retries=retries,
        note=note,
        debug=debug,
        timeout=timeout,
    )

    if (type(conn_params.get('ontology')) == list and len(conn_params.get('ontology')) > 1) or ',' in conn_params.get('ontology'):
        conn_params = determine_concept_res.get('conn_params')

    concept = determine_concept_res.get('concept')
    identify_concept_reason = determine_concept_res.get('identify_concept_reason', None)
    schema = determine_concept_res.get('schema')
    concept_metadata = determine_concept_res.get('concept_metadata')
    usage_metadata.update(determine_concept_res.get('usage_metadata', {}))

    if not concept:
        raise Exception("No relevant concept found for the query.")

    generate_sql_prompt = get_generate_sql_prompt_template(conn_params)
    sql_query = None
    generate_sql_reason = None
    is_sql_valid = True  # Assume valid by default; set to False only if validation fails
    error = ''

    try:
        result = _generate_sql_with_llm(
            question=question,
            llm=llm,
            generate_sql_prompt=generate_sql_prompt,
            current_context=_build_sql_generation_context(
                conn_params=conn_params,
                schema=schema,
                concept=concept,
                concept_metadata=concept_metadata,
                graph_depth=graph_depth,
                include_tags=include_tags,
                exclude_properties=exclude_properties,
                db_is_case_sensitive=db_is_case_sensitive,
                max_limit=max_limit),
            note=note,
            timeout=timeout,
            debug=debug,
        )
        
        usage_metadata['generate_sql'] = {
            "approximate": result['apx_token_count'],
            **result['usage_metadata'],
        }
        if debug and 'p_hash' in result:
            usage_metadata['generate_sql']["p_hash"] = result['p_hash']
        
        sql_query = result['sql']
        generate_sql_reason = result.get('generate_sql_reason', None)
        error = result['error']

        if error:
            raise Exception(error)
        
        if enable_reasoning and sql_query is not None:
            sql_query, graph_depth, generate_sql_reason = handle_generate_sql_reasoning(
                sql_query=sql_query,
                question=question,
                llm=llm,
                conn_params=conn_params,
                schema=schema,
                concept=concept,
                concept_metadata=concept_metadata,
                include_tags=include_tags,
                exclude_properties=exclude_properties,
                db_is_case_sensitive=db_is_case_sensitive,
                max_limit=max_limit,
                reasoning_steps=reasoning_steps,
                note=note,
                graph_depth=graph_depth,
                usage_metadata=usage_metadata,
                timeout=timeout,
                debug=debug,
                previous_token_count=result['apx_token_count']
            )

        if should_validate_sql or enable_reasoning:
            # Validate & regenerate only once if reasoning enabled and validation is disabled
            validate_retries = 1 if not should_validate_sql else retries
            is_sql_valid, error, sql_query = handle_validate_generate_sql(
                sql_query=sql_query,
                question=question,
                llm=llm,
                conn_params=conn_params,
                generate_sql_prompt=generate_sql_prompt,
                schema=schema,
                concept=concept,
                concept_metadata=concept_metadata,
                include_tags=include_tags,
                exclude_properties=exclude_properties,
                db_is_case_sensitive=db_is_case_sensitive,
                max_limit=max_limit,
                graph_depth=graph_depth,
                retries=validate_retries,
                timeout=timeout,
                debug=debug,
                usage_metadata=usage_metadata,
            )
    except TimeoutError as e:
        error = f"LLM call timed out: {str(e)}"
        raise Exception(error)
    except Exception as e:
        error = f"LLM call failed: {str(e)}"
        raise Exception(error)
    
    return {
        "sql": sql_query,
        "concept": concept,
        "schema": schema,
        "error": error if not is_sql_valid else None,
        "is_sql_valid": is_sql_valid if should_validate_sql else None,
        "identify_concept_reason": identify_concept_reason,
        "generate_sql_reason": generate_sql_reason,
        "reasoning_status": reasoning_status,
        "usage_metadata": usage_metadata,
        "ontology": conn_params.get('ontology'),
        "conn_params": conn_params
    }


def answer_question(
    question: str,
    llm: LLM,
    conn_params: dict,
    results: str,
    sql: Optional[str] = None,
    timeout: Optional[int] = None,
    note: Optional[str] = '',
    debug: Optional[bool] = False,
) -> dict[str, Any]:
    # Use config default timeout if none provided
    if timeout is None:
        timeout = config.llm_timeout

    qa_prompt = get_qa_prompt_template(conn_params)

    prompt = qa_prompt.format_messages(
        question=question,
        formatted_rows=results,
        additional_context=f"SQL QUERY:\n{sql}\n\n" if sql else "",
        note=note,
    )
    
    apx_token_count = _calculate_token_count(llm, prompt)

    if "snowflake" in llm._llm_type:
        _clean_snowflake_prompt(prompt)
    
    try:
        response = _call_llm_with_timeout(llm, prompt, timeout=timeout)
    except TimeoutError as e:
        raise TimeoutError(f"LLM call timed out while answering question: {str(e)}")
    except Exception as e:
        raise Exception(f"LLM call failed while answering question: {str(e)}")

    if hasattr(response, "content"):
        response_text = response.content
    elif isinstance(response, str):
        response_text = response
    else:
        raise ValueError("Unexpected response format from LLM.")
    
    usage_metadata = {
        "answer_question": {
            "approximate": apx_token_count,
            **_extract_usage_metadata(response),
        },
    }
    if debug:
        usage_metadata["answer_question"]["p_hash"] = encrypt_prompt(prompt)

    return {
        "answer": response_text,
        "usage_metadata": usage_metadata,
    }

