import os
import sys
from pathlib import Path

# Ensure Backend directory is in sys.path and root .env is loaded
BACKEND_DIR = Path(__file__).resolve().parent.parent
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))

import config
from pydantic import BaseModel
from litellm import completion
from database.database import execute_query

# Simplified schema reference for the LLM context
DB_SCHEMA = """
CREATE TABLE customers (id INTEGER PRIMARY KEY, first_name TEXT, last_name TEXT, email TEXT, created_at TIMESTAMP);
CREATE TABLE orders (id INTEGER PRIMARY KEY, customer_id INTEGER, product_name TEXT, amount REAL, status TEXT, order_date TIMESTAMP, FOREIGN KEY (customer_id) REFERENCES customers(id));
"""

class QueryToolParams(BaseModel):
    natural_language_query: str

def run_query_tool(params: dict) -> str:
    """Translates NL to SQL, validates it, and executes it against the database."""
    try:
        tool_config = QueryToolParams(**params)
    except Exception as e:
        return f"Error parsing parameters: {e}"

    prompt = f"""You are an expert SQL assistant. Convert the following natural language query into a raw SQL SELECT statement.
Schema:
{DB_SCHEMA}

Query: {tool_config.natural_language_query}

Return ONLY the SQL string, without markdown formatting or code blocks.
"""
    try:
        model_name = config.get_llm_model()
        api_key = config.get_api_key()
        
        completion_kwargs = {
            "model": model_name,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0
        }
        if api_key:
            completion_kwargs["api_key"] = api_key

        response = completion(**completion_kwargs)
        message = response.choices[0].message
        sql_query = (message.content or getattr(message, "reasoning_content", "") or "").strip()
    except Exception as e:
        return f"Query Tool Error: {str(e)}"
        
    try:
        import re
        # If output contains markdown code block, extract it
        code_block_match = re.search(r"```(?:sql)?\s*([\s\S]*?)\s*```", sql_query, re.IGNORECASE)
        if code_block_match:
            sql_query = code_block_match.group(1).strip()
        elif not sql_query.upper().startswith("SELECT"):
            # If wrapped in reasoning text, extract the SELECT statement
            select_match = re.search(r"(SELECT\s+[\s\S]+?)(?:;|\n\n|$)", sql_query, re.IGNORECASE)
            if select_match:
                sql_query = select_match.group(1).strip()
        if sql_query.startswith("```sql"):
            sql_query = sql_query.replace("```sql", "", 1)
            if sql_query.endswith("```"):
                sql_query = sql_query[::-1].replace("```", "", 1)[::-1]
        elif sql_query.startswith("```"):
            sql_query = sql_query.replace("```", "", 1)
            if sql_query.endswith("```"):
                sql_query = sql_query[::-1].replace("```", "", 1)[::-1]
        
        sql_query = sql_query.strip()
        
        # Simple read-only check since safety module is removed
        if not sql_query.upper().startswith("SELECT"):
            return "Safety Violation: Only SELECT queries are allowed."
            
        if "LIMIT" not in sql_query.upper():
            sql_query = sql_query.rstrip(";") + " LIMIT 50;"
            
        rows = execute_query(sql_query)
        if not rows:
            return f"Executed `{sql_query}`: No results found."
        
        return f"Executed `{sql_query}`:\n" + str(rows)

    except Exception as e:
        return f"Query Tool Error: {str(e)}"