"""
数据源连接器服务
"""
from typing import Dict, Any, List, Optional
import logging
import json

logger = logging.getLogger(__name__)


class DataSourceConnector:
    """数据源连接器基类"""
    
    def __init__(self, source_type: str, config: Dict[str, Any]):
        """
        初始化数据源连接器
        
        Args:
            source_type: 数据源类型
            config: 连接配置
        """
        self.source_type = source_type
        self.config = config
    
    def test_connection(self) -> Dict[str, Any]:
        """
        测试连接
        
        Returns:
            连接测试结果
        """
        raise NotImplementedError("子类必须实现test_connection方法")
    
    def query(self, query_params: Dict[str, Any]) -> Any:
        """
        查询数据
        
        Args:
            query_params: 查询参数
            
        Returns:
            查询结果
        """
        raise NotImplementedError("子类必须实现query方法")


class MySQLConnector(DataSourceConnector):
    """MySQL连接器"""
    
    def test_connection(self) -> Dict[str, Any]:
        try:
            import pymysql
            connection = pymysql.connect(
                host=self.config.get('host'),
                port=self.config.get('port', 3306),
                user=self.config.get('user'),
                password=self.config.get('password'),
                database=self.config.get('database'),
                connect_timeout=5
            )
            connection.close()
            return {"status": "success", "message": "连接成功"}
        except Exception as e:
            raise Exception(f"MySQL连接失败: {str(e)}")
    
    def query(self, query_params: Dict[str, Any]) -> List[Dict[str, Any]]:
        try:
            import pymysql
            sql = query_params.get('sql')
            if not sql:
                raise ValueError("缺少SQL查询语句")
            
            connection = pymysql.connect(
                host=self.config.get('host'),
                port=self.config.get('port', 3306),
                user=self.config.get('user'),
                password=self.config.get('password'),
                database=self.config.get('database')
            )
            
            try:
                with connection.cursor(pymysql.cursors.DictCursor) as cursor:
                    cursor.execute(sql)
                    result = cursor.fetchall()
                    return result
            finally:
                connection.close()
        except Exception as e:
            raise Exception(f"MySQL查询失败: {str(e)}")


class PostgreSQLConnector(DataSourceConnector):
    """PostgreSQL连接器"""
    
    def test_connection(self) -> Dict[str, Any]:
        try:
            import psycopg2
            connection = psycopg2.connect(
                host=self.config.get('host'),
                port=self.config.get('port', 5432),
                user=self.config.get('user'),
                password=self.config.get('password'),
                database=self.config.get('database'),
                connect_timeout=5
            )
            connection.close()
            return {"status": "success", "message": "连接成功"}
        except Exception as e:
            raise Exception(f"PostgreSQL连接失败: {str(e)}")
    
    def query(self, query_params: Dict[str, Any]) -> List[Dict[str, Any]]:
        try:
            import psycopg2
            import psycopg2.extras
            sql = query_params.get('sql')
            if not sql:
                raise ValueError("缺少SQL查询语句")
            
            connection = psycopg2.connect(
                host=self.config.get('host'),
                port=self.config.get('port', 5432),
                user=self.config.get('user'),
                password=self.config.get('password'),
                database=self.config.get('database')
            )
            
            try:
                with connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                    cursor.execute(sql)
                    result = cursor.fetchall()
                    return [dict(row) for row in result]
            finally:
                connection.close()
        except Exception as e:
            raise Exception(f"PostgreSQL查询失败: {str(e)}")


class APIConnector(DataSourceConnector):
    """API连接器"""
    
    def test_connection(self) -> Dict[str, Any]:
        try:
            import httpx
            url = self.config.get('base_url')
            if not url:
                raise ValueError("缺少base_url配置")
            
            headers = self.config.get('headers', {})
            timeout = self.config.get('timeout', 10)
            
            response = httpx.get(url, headers=headers, timeout=timeout)
            response.raise_for_status()
            
            return {"status": "success", "message": "连接成功", "status_code": response.status_code}
        except Exception as e:
            raise Exception(f"API连接失败: {str(e)}")
    
    def query(self, query_params: Dict[str, Any]) -> Any:
        try:
            import httpx
            method = query_params.get('method', 'GET').upper()
            endpoint = query_params.get('endpoint', '')
            params = query_params.get('params', {})
            data = query_params.get('data', {})
            headers = self.config.get('headers', {})
            timeout = self.config.get('timeout', 10)
            
            base_url = self.config.get('base_url', '').rstrip('/')
            url = f"{base_url}/{endpoint.lstrip('/')}"
            
            if method == 'GET':
                response = httpx.get(url, params=params, headers=headers, timeout=timeout)
            elif method == 'POST':
                response = httpx.post(url, json=data, headers=headers, timeout=timeout)
            elif method == 'PUT':
                response = httpx.put(url, json=data, headers=headers, timeout=timeout)
            elif method == 'DELETE':
                response = httpx.delete(url, headers=headers, timeout=timeout)
            else:
                raise ValueError(f"不支持的HTTP方法: {method}")
            
            response.raise_for_status()
            return response.json() if response.content else {}
        except Exception as e:
            raise Exception(f"API查询失败: {str(e)}")


class JSONFileConnector(DataSourceConnector):
    """JSON文件连接器"""
    
    def test_connection(self) -> Dict[str, Any]:
        try:
            import os
            file_path = self.config.get('file_path')
            if not file_path:
                raise ValueError("缺少file_path配置")
            
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件不存在: {file_path}")
            
            return {"status": "success", "message": "文件存在"}
        except Exception as e:
            raise Exception(f"JSON文件连接失败: {str(e)}")
    
    def query(self, query_params: Dict[str, Any]) -> Any:
        try:
            import json
            import os
            file_path = self.config.get('file_path')
            
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件不存在: {file_path}")
            
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # 支持简单的查询过滤
            filter_path = query_params.get('path')
            if filter_path:
                # 支持JSONPath风格的路径查询
                parts = filter_path.split('.')
                result = data
                for part in parts:
                    if isinstance(result, dict):
                        result = result.get(part)
                    elif isinstance(result, list):
                        try:
                            index = int(part)
                            result = result[index]
                        except (ValueError, IndexError):
                            return None
                    else:
                        return None
                return result
            
            return data
        except Exception as e:
            raise Exception(f"JSON文件查询失败: {str(e)}")


# 连接器工厂
_connector_classes = {
    'mysql': MySQLConnector,
    'postgresql': PostgreSQLConnector,
    'api': APIConnector,
    'json': JSONFileConnector,
}


def create_connector(source_type: str, config: Dict[str, Any]):
    """
    创建数据源连接器
    
    Args:
        source_type: 数据源类型
        config: 连接配置
        
    Returns:
        数据源连接器实例
    """
    connector_class = _connector_classes.get(source_type)
    if not connector_class:
        raise ValueError(f"不支持的数据源类型: {source_type}")
    
    return connector_class(source_type, config)


# 为了兼容API，创建一个统一的DataSourceConnector包装类
class DataSourceConnectorWrapper:
    """统一的数据源连接器包装类（用于API调用）"""
    
    def __init__(self, source_type: str, config: Dict[str, Any]):
        self.connector = create_connector(source_type, config)
        self.source_type = source_type
        self.config = config
    
    def test_connection(self) -> Dict[str, Any]:
        return self.connector.test_connection()
    
    def query(self, query_params: Dict[str, Any]) -> Any:
        return self.connector.query(query_params)


# 导出时使用包装类，这样API可以统一使用DataSourceConnector
# 但实际返回的是具体的连接器实现
