"""
Interactive SQL client for querying slowlog data.
"""
from slowlogs.ruler import DeltaLakeStore
from app.logcli.utils.logger import log_info, log_error
import os
import sys
import pandas as pd
import requests
import readline
import atexit
from typing import Optional
import os.path
import time
from pathlib import Path

# 历史文件路径
HISTORY_FILE = os.path.expanduser('~/.slowlog_sql_history')


def save_history():
    readline.write_history_file(HISTORY_FILE)


# 加载历史记录
if os.path.exists(HISTORY_FILE):
    readline.read_history_file(HISTORY_FILE)

# 设置历史记录长度
readline.set_history_length(1000)

# 注册退出时保存历史记录
atexit.register(save_history)

# Add the project root to the Python path
project_root = str(Path(__file__).parent.parent)
if project_root not in sys.path:
    sys.path.append(project_root)


class SlowlogSQLClient:
    """Interactive SQL client for querying slowlog data."""

    def __init__(self, vendor: str, region: str, tenant_id: str, cluster_id: str):
        """Initialize the SQL client with cluster information."""
        self.vendor = vendor
        self.region = region
        self.tenant_id = tenant_id
        self.cluster_id = cluster_id

        # Initialize AWS credentials
        self._setup_aws_credentials()

        # Get auth token
        from auth.client_auth import get_credentials
        creds = get_credentials()
        auth0_token = creds.get('auth0_token')

        # Query cluster info using the same logic as detector
        url = "http://www.gs.us-west-2.aws.observability.tidbcloud.com/api/v1/clusters/"
        url += f"?bizType=tidbcloud&vendor={vendor}&region={region}&tenantID={tenant_id}&applicationID={cluster_id}"

        headers = {'authorization': f'Bearer {auth0_token}'}

        log_info(f"Querying cluster info from: {url}")
        resp = requests.get(url, headers=headers)
        resp.raise_for_status()

        clusters = resp.json()["data"]["clusters"]
        if not clusters:
            raise ValueError(
                f"No cluster found with ID {cluster_id} in tenant {tenant_id}")

        # Filter out system clusters using the same logic as detector
        exclude_application_ids = ['k8s-infra', 'base', 'seed']
        exclude_project_ids = ['seed', "base"]

        filtered_clusters = [
            c for c in clusters
            if not (c.get('applicationID') in exclude_application_ids 
                    or c.get('projectID') in exclude_project_ids)
        ]

        if not filtered_clusters:
            raise ValueError(
                f"No valid clusters found after filtering for ID {cluster_id}")

        cluster_info = filtered_clusters[0]

        # Initialize cluster with the correct parameters
        # Initialize store
        self.store = DeltaLakeStore(
            vendor=vendor,
            region=cluster_info['region'],
            tenant_id=cluster_info['tenantID'],
            o11y_id=cluster_info['id'])

        log_info("Initialized cluster: ")
        log_info(f"- Region: {self.store.region}")
        log_info(f"- Tenant ID: {self.store.tenant_id}")
        log_info(f"- Observability ID: {self.store.o11y_id}")
        log_info(f"- Application ID: {cluster_info.get('applicationID')}")
        log_info(f"- Project ID: {cluster_info.get('projectID')}")

        # Initialize connection and register tables
        self.connection = self.store.db  # Use the store's connection
        # self._register_tables()

    def _setup_aws_credentials(self):
        """Set up AWS credentials using auth.client_auth."""
        try:
            from auth.client_auth import get_credentials
            creds = get_credentials()

            # Set AWS environment variables
            os.environ['AWS_REGION'] = self.region
            os.environ['AWS_ACCESS_KEY_ID'] = creds['access_key_id']
            os.environ['AWS_SECRET_ACCESS_KEY'] = creds['secret_access_key']
            os.environ['AWS_SESSION_TOKEN'] = creds.get('session_token', '')

            log_info("Successfully set up AWS credentials")
        except Exception as e:
            log_error(f"Failed to set up AWS credentials: {str(e)}")
            log_error(
                "Please make sure you have valid AWS credentials configured")
            raise

    def _get_connection(self):
        """Get the database connection."""
        return self.connection

    def execute_query(self, query: str, limit: int = 100) -> Optional[pd.DataFrame]:
        """
        Execute a SQL query and return the results as a pandas DataFrame.

        Args:
            query: SQL query to execute
            limit: Maximum number of rows to return (0 for no limit)

        Returns:
            DataFrame with query results or None if an error occurred
        """
        try:
            # Add LIMIT if not already present and limit > 0
            if limit > 0 and 'LIMIT' not in query.upper() and 'COPY' not in query.upper():
                query = f"{query.rstrip(';')} LIMIT {limit};"
            start_time = time.time()
            conn = self._get_connection()
            result = conn.execute(query).fetchdf()
            end_time = time.time()
            if not result.empty:
                print(f"\n{'-' * 80}")
                print(
                    f"Query returned {len(result)} row{'s' if len(result) != 1 else ''}")
                print(f"Query took {end_time - start_time:.2f} seconds")
                print("-" * 80)
                with pd.option_context('display.max_columns', None, 'display.width', None, 'max_colwidth', 100):
                    print(result)
                print("-" * 80)
            else:
                print("No results returned.")

            return result

        except Exception as e:
            log_error(f"Error executing query: {e}")
            return None

    def interactive(self):
        """Start an interactive SQL shell."""
        print("\n" + "=" * 60)
        print("  Slowlog Interactive SQL Shell")
        print("  Type 'tables' to list available tables")
        print("  Type 'schema [table]' to view table schema")
        print("  Type 'exit' or 'quit' to exit")
        print("  some common sql example:")
        print("    SELECT id, to_timestamp(time), query_time, query FROM slowlogs WHERE time>epoch(TIMESTAMP '2025-06-01 05:50:00') ORDER BY query_time DESC;")
        print("    SELECT to_timestamp((time / 600) * 600) AS time_bucket, COUNT(*) AS cnt FROM slowlogs GROUP BY time_bucket ORDER BY time_bucket;")
        print("  Use up/down arrow keys to navigate command history")
        print("=" * 60 + "\n")

        # 初始化历史记录
        history = []
        history_index = 0

        def pre_input():
            nonlocal history_index
            if history_index > 0:
                history_index -= 1
                sys.stdout.write("\r" + " " * 100 + "\r")  # 清除当前行
                return history[history_index]
            return ""

        def next_input():
            nonlocal history_index
            if history_index < len(history) - 1:
                history_index += 1
                sys.stdout.write("\r" + " " * 100 + "\r")  # 清除当前行
                return history[history_index]
            elif history_index == len(history) - 1:
                history_index += 1
                sys.stdout.write("\r" + " " * 100 + "\r")  # 清除当前行
                return ""
            return ""

        # 设置readline自动补全
        def complete(text, state):
            commands = ['tables', 'schema', 'exit', 'quit', 'SELECT',
                        'FROM', 'WHERE', 'GROUP BY', 'ORDER BY', 'LIMIT']
            matches = [cmd for cmd in commands if cmd.startswith(text.upper())]
            return matches[state] if state < len(matches) else None

        readline.parse_and_bind('tab: complete')
        readline.set_completer(complete)

        while True:
            try:
                # Read input
                try:
                    query = input("sql> ").strip()

                    # 保存非空查询到历史记录
                    if query and (not history or history[-1] != query):
                        history.append(query)
                        history_index = len(history)
                        readline.add_history(query)  # 添加到readline历史
                        save_history()  # 保存历史记录到文件
                except EOFError:
                    print("\nGoodbye!")
                    break

                if not query:
                    continue

                # Handle special commands
                if query.lower() in ('exit', 'quit'):
                    print("Goodbye!")
                    break

                if query.lower() == 'tables':
                    self._show_tables()
                    continue

                if query.lower().startswith('schema'):
                    parts = query.split()
                    if len(parts) == 2:
                        self._show_schema(parts[1])
                    else:
                        print("Usage: schema <table_name>")
                    continue

                # Execute SQL query
                self.execute_query(query)

            except KeyboardInterrupt:
                print("\nType 'exit' or 'quit' to exit")

            except Exception as e:
                log_error(f"Error: {e}")

    def _show_tables(self):
        """Show available tables and their row counts."""
        try:
            conn = self._get_connection()
            tables = conn.execute("SHOW TABLES;").fetchall()

            if not tables:
                print("No tables found.")
                return

            print("\nAvailable tables:")
            print("-" * 60)
            for table in tables:
                print(f"{table[0]}")
            print("-" * 60)

        except Exception as e:
            log_error(f"Error listing tables: {e}")

    def _show_schema(self, table_name: str):
        """Show schema for a specific table."""
        try:
            conn = self._get_connection()

            # First check if table exists
            tables = [t[0].lower()
                      for t in conn.execute("SHOW TABLES;").fetchall()]
            if table_name.lower() not in tables:
                print(
                    f"Table '{table_name}' not found. Available tables: {', '.join(tables)}")
                return

            # Get column info
            schema = conn.execute(f"DESCRIBE {table_name};").fetchdf()

            # Configure pandas to show all columns and rows
            pd.set_option('display.max_columns', None)
            pd.set_option('display.width', None)
            pd.set_option('display.max_colwidth', None)

            print(f"\nSchema for table '{table_name}':")
            print("-" * 100)  # Wider separator for better visibility
            with pd.option_context('display.max_rows', None, 'display.max_columns', None, 'display.width', None, 'max_colwidth', None):
                print(schema)
            print("-" * 100)

        except Exception as e:
            log_error(f"Error showing schema: {e}")

    def __del__(self):
        """Clean up resources."""
        if self.connection:
            self.connection.close()