from sqlalchemy import create_engine, inspect
from typing import List, Dict, Any
from .config import DBConfig

class TableMetadata:
    def __init__(self, name: str, columns: List[Dict[str, Any]], primary_keys: List[str]):
        self.name = name
        self.columns = columns
        self.primary_keys = primary_keys
        
        # 转换表名为类名 (例如: users -> User, order_items -> OrderItem)
        self.class_name = "".join([part.capitalize() for part in name.split("_")])
        if self.class_name.endswith("s") and len(self.class_name) > 1:
            self.class_name = self.class_name[:-1]

def inspect_db(config: DBConfig) -> List[TableMetadata]:
    engine = create_engine(config.connection_url)
    inspector = inspect(engine)
    
    schema = config.schema_name if config.schema_name else None
    table_names = config.tables if config.tables else inspector.get_table_names(schema=schema)
    metadata_list = []
    
    for table_name in table_names:
        columns = inspector.get_columns(table_name, schema=schema)
        pk_constraint = inspector.get_pk_constraint(table_name, schema=schema)
        primary_keys = pk_constraint.get("constrained_columns", [])
        
        # 处理列类型，将其映射为 Python 类型字符串
        processed_columns = []
        for col in columns:
            col_type_obj = col["type"]
            col_type_str = col_type_obj.__class__.__name__
            
            # 某些类型可能不在模板的简单 import 列表中，确保基础类型正确
            if col_type_str == "INTEGER": col_type_str = "Integer"
            if col_type_str == "BIGINT": col_type_str = "BigInteger"
            if col_type_str == "DATE": col_type_str = "Date"
            if col_type_str == "VARCHAR": col_type_str = "String"
            if col_type_str == "DATETIME": col_type_str = "DateTime"
            if col_type_str == "TIMESTAMP": col_type_str = "DateTime"
            if col_type_str == "FLOAT": col_type_str = "Float"
            if col_type_str == "BOOLEAN": col_type_str = "Boolean"
            if col_type_str == "TEXT": col_type_str = "Text"
            if col_type_str == "NUMERIC": col_type_str = "Numeric"
            if col_type_str == "DECIMAL": col_type_str = "Numeric"
            if col_type_str == "UUID": col_type_str = "String"  # 简单处理为 String
            if col_type_str == "JSON" or col_type_str == "JSONB": col_type_str = "JSON"
            
            python_type = "str"
            col_type_lower = str(col_type_obj).lower()
            if "int" in col_type_lower:
                python_type = "int"
            elif "float" in col_type_lower or "decimal" in col_type_lower or "numeric" in col_type_lower:
                python_type = "float"
            elif "bool" in col_type_lower:
                python_type = "bool"
            elif "datetime" in col_type_lower or "timestamp" in col_type_lower:
                python_type = "datetime"
            
            processed_columns.append({
                "name": col["name"],
                "type": col_type_str,
                "python_type": python_type,
                "nullable": col.get("nullable", True),
                "default": col.get("default"),
                "is_pk": col["name"] in primary_keys
            })
            
        metadata_list.append(TableMetadata(table_name, processed_columns, primary_keys))
        
    return metadata_list
