import sys
import os
import json
from pathlib import Path
from typing import List, Optional, Dict, Any
from pydantic import BaseModel, Field

# Add project root to sys.path to allow importing nanobot
PROJECT_ROOT = Path(__file__).resolve().parents[3]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.append(str(PROJECT_ROOT))

from nanobot.providers.litellm_provider import LiteLLMProvider
from app.connectors.postgres import postgres_connector
from app.connectors.clickhouse import clickhouse_connector
from app.api.llm import _load_data as load_llm_config

class NL2SQLRequest(BaseModel):
    query: str = Field(..., description="User's natural language query")
    source: str = Field(..., description="Data source to query (postgres, clickhouse)")

class NL2SQLResponse(BaseModel):
    sql: str
    result: List[Dict[str, Any]]
    error: Optional[str] = None

async def process_nl2sql(request: NL2SQLRequest) -> NL2SQLResponse:
    # 1. Get the connector and schema
    connector = None
    if request.source == "postgres":
        connector = postgres_connector
    elif request.source == "clickhouse":
        connector = clickhouse_connector
    else:
        return NL2SQLResponse(sql="", result=[], error=f"Unsupported data source: {request.source}")

    if not connector.test_connection():
         return NL2SQLResponse(sql="", result=[], error=f"Failed to connect to {request.source}")

    schema = connector.get_schema()
    schema_str = json.dumps(schema, indent=2)

    # 2. Get the active LLM config
    llm_configs = load_llm_config()
    active_config = next((c for c in llm_configs if c.get("is_active")), None)
    
    if not active_config:
        return NL2SQLResponse(sql="", result=[], error="No active LLM configuration found")

    # 3. Initialize Provider
    try:
        provider = LiteLLMProvider(
            api_key=active_config.get("api_key"),
            api_base=active_config.get("api_base"),
            default_model=active_config.get("model"),
            extra_headers=active_config.get("extra_headers")
        )
    except Exception as e:
        return NL2SQLResponse(sql="", result=[], error=f"Failed to initialize LLM provider: {e}")

    # 4. Construct Prompt
    prompt = f"""You are an expert SQL generator. 
Given the following database schema for a {request.source} database:
{schema_str}

Write a SQL query to answer the following question:
"{request.query}"

Return ONLY the SQL query. Do not include any markdown formatting, explanations, or code blocks. Just the raw SQL string.
"""

    # 5. Call LLM
    try:
        # provider.complete returns a string
        response = await provider.complete(prompt)
        sql_query = response.strip()
        # Remove potential markdown code blocks if the LLM ignores instructions
        if sql_query.startswith("```sql"):
            sql_query = sql_query[6:]
        if sql_query.startswith("```"):
            sql_query = sql_query[3:]
        if sql_query.endswith("```"):
            sql_query = sql_query[:-3]
        sql_query = sql_query.strip()
    except Exception as e:
        return NL2SQLResponse(sql="", result=[], error=f"LLM generation failed: {e}")

    # 6. Execute SQL
    try:
        results = connector.execute_query(sql_query)
        # Convert results to list of dicts if not already (Postgres returns list of dicts, ClickHouse returns list of tuples)
        formatted_results = []
        if request.source == "postgres":
             formatted_results = results
        elif request.source == "clickhouse":
            # ClickHouse returns list of tuples, we need column names
            # But execute_query in ClickHouseConnector just returns raw results from client.execute
            # client.execute(query, with_column_types=True) might be better but let's stick to simple for now
            # Actually, without column names it's hard to format as dict.
            # Let's assume we can just return the raw tuples for now or try to fetch column names.
            # For now, let's just return as list of lists/tuples if it's not a dict
            formatted_results = [list(row) for row in results]

        return NL2SQLResponse(sql=sql_query, result=formatted_results)
    except Exception as e:
        return NL2SQLResponse(sql=sql_query, result=[], error=f"SQL execution failed: {e}")