import os
import logging
from contextlib import asynccontextmanager
from collections.abc import AsyncGenerator
from typing import Dict, Any, Optional, List, Union

from pydantic import Field, BaseModel, ConfigDict

from mcp.server.fastmcp import FastMCP

from alibabacloud_dms_enterprise20181101.client import Client as dms_enterprise20181101Client
from alibabacloud_tea_openapi import models as open_api_models
from alibabacloud_dms_enterprise20181101 import models as dms_enterprise_20181101_models

# --- Global Logger ---
logger = logging.getLogger(__name__)


# --- Pydantic Models ---
class MyBaseModel(BaseModel):
    model_config = ConfigDict(json_dumps_params={'ensure_ascii': False})


class AskDatabaseResult(MyBaseModel):
    executed_sql: str = Field(description="The SQL query that was executed.")
    execution_result: str = Field(description="The result of the SQL query execution.")


class InstanceInfo(MyBaseModel):
    InstanceId: Any = Field(description="Unique instance identifier in DMS", default=None)
    Host: Any = Field(description="The hostname of the database instance", default=None)
    Port: Any = Field(description="The connection port number", default=None)


class InstanceDetail(MyBaseModel):
    InstanceId: Any = Field(description="Unique instance identifier in DMS", default=None)
    State: Any = Field(description="Current operational status", default=None)
    InstanceType: Any = Field(description="Database Engine type", default=None)
    InstanceAlias: Any = Field(description="Instance alias in DMS", default=None)
    EnvType: Any = Field(description="The environment type of the instance (e.g., production, development, etc.)",
                         default=None)
    Host: Any = Field(description="The hostname of the database instance", default=None)
    Port: Any = Field(description="The connection port number", default=None)
    InstanceSource: Any = Field(description="The instance source (e.g., RDS, VPC_IDC, ECS_OWN, PUBLIC_OWN etc.)",
                                default=None)
    InstanceResourceId: Any = Field(
        description="Resource ID of the instance from RDS",
        default=None)


class DatabaseInfo(MyBaseModel):
    DatabaseId: Any = Field(description="Unique database identifier in DMS")
    Host: Any = Field(description="Hostname or IP address of the database instance")
    Port: Any = Field(description="Connection port number")
    DbType: Any = Field(description="Database Engine type")
    SchemaName: Any = Field(description="Name of the database schema")


class DatabaseDetail(MyBaseModel):
    DatabaseId: Any = Field(description="Unique database identifier in DMS", default=None)
    SchemaName: Any = Field(description="Name of the database schema", default=None)
    DbType: Any = Field(description="Database Engine type", default=None)
    InstanceAlias: Any = Field(description="Instance alias in DMS", default=None)
    InstanceId: Any = Field(description="Instance identifier in DMS", default=None)
    State: Any = Field(description="Current operational status", default=None)


class Column(MyBaseModel):
    ColumnName: Any = Field(description="Name of the column")
    ColumnType: Any = Field(description="Full SQL type declaration (e.g., 'varchar(32)', 'bigint(20)')")
    AutoIncrement: Any = Field(description="Whether the column is an auto-increment field")
    Description: Any = Field(description="Column comment/description text")
    Nullable: Any = Field(description="Whether NULL values are allowed")


class Index(MyBaseModel):
    IndexColumns: Any = Field(description="List of column names included in the index")
    IndexName: Any = Field(description="Name of the index")
    IndexType: Any = Field(description="Type of index ('Primary', 'Unique', etc.)")
    Unique: Any = Field(description="Whether the index enforces uniqueness")


class TableDetail(MyBaseModel):
    ColumnList: Any = Field(description="List of column metadata", default=None)
    IndexList: Any = Field(description="List of index metadata", default=None)


class ResultSet(MyBaseModel):
    ColumnNames: List[str] = Field(description="Ordered list of column names")
    RowCount: int = Field(description="Number of rows returned")
    Rows: List[Dict[str, Any]] = Field(description="List of rows, where each row is a dictionary of column_name: value")
    MarkdownTable: Optional[str] = Field(default=None, description="Data formatted as a Markdown table string")
    Success: bool = Field(description="Whether this result set was successfully retrieved")
    Message: str = Field(description="Additional message returned")


class ExecuteScriptResult(MyBaseModel):
    RequestId: str = Field(description="Unique request identifier")
    Results: List[ResultSet] = Field(description="List of result sets from executed script")
    Success: bool = Field(description="Overall operation success status")

    def __str__(self) -> str:
        if self.Success and self.Results:
            first_result = self.Results[0]
            if first_result.Success and first_result.MarkdownTable:
                return first_result.MarkdownTable
            elif not first_result.Success:
                return first_result.Message
            else:
                return "Result data is not available in Markdown format."
        elif not self.Success:
            return "Script execution failed."
        else:
            return "Script executed successfully, but no results were returned."


class SqlResult(MyBaseModel):
    sql: Optional[str] = Field(description="The generated SQL query")


DATABASE_ID_DESCRIPTION = (
    "If you don't know the databaseId, first use getDatabase or searchDatabase to retrieve it.\n"
    "(1) If you have the exact host, port, and database name, use getDatabase.\n"
    "(2) If you only know the database name, use searchDatabase.\n"
    "(3) If you don't know any information, ask the user to provide the necessary details.\n"
    "Note: searchDatabase may return multiple databases. In this case, let the user choose which one to use."
)


# --- Aliyun Client Creation ---
def create_client() -> dms_enterprise20181101Client:
    config = open_api_models.Config(
        access_key_id=os.getenv('ALIBABA_CLOUD_ACCESS_KEY_ID'),
        access_key_secret=os.getenv('ALIBABA_CLOUD_ACCESS_KEY_SECRET'),
        security_token=os.getenv('ALIBABA_CLOUD_SECURITY_TOKEN'),
        read_timeout=60 * 1000  # 设置读取超时时间为60秒
    )
    config.endpoint = f'dms-enterprise.cn-hangzhou.aliyuncs.com'
    config.user_agent = "dms-mcp"
    return dms_enterprise20181101Client(config)


async def add_instance(
        db_user: str = Field(description="The username used to connect to the database"),
        db_password: str = Field(description="The password used to connect to the database"),
        instance_resource_id: Optional[str] = Field(default=None, description="Aliyun instance resource ID"),
        host: Optional[str] = Field(default=None, description="The hostname of the database instance"),
        port: Optional[str] = Field(default=None, description="The connection port number"),
        region: Optional[str] = Field(default=None, description="The region (e.g., 'cn-hangzhou')")
) -> InstanceInfo:
    if not db_user or not isinstance(db_user, str):
        raise ValueError("db_user must be a non-empty string")
    if not db_password or not isinstance(db_password, str):
        raise ValueError("db_password must be a non-empty string")
    client = create_client()
    req = dms_enterprise_20181101_models.SimplyAddInstanceRequest(database_user=db_user, database_password=db_password)
    if host:
        req.host = host
    if port:
        req.port = port
    if instance_resource_id:
        req.instance_id = instance_resource_id
    if region:
        req.region = region
    if mcp.state.real_login_uid:
        req.real_login_user_uid = mcp.state.real_login_uid
    try:
        resp = client.simply_add_instance(req)

        if resp and resp.body:
            body_dict = resp.body.to_map()
            return InstanceInfo(**(body_dict['Instance'] if 'Instance' in body_dict else body_dict))
        else:
            return InstanceInfo()

    except Exception as e:
        logger.error(f"Error in add_instance: {e}")
        raise


async def get_instance(
        host: str = Field(description="The hostname of the database instance"),
        port: str = Field(description="The connection port number"),
        sid: Optional[str] = Field(default=None, description="Required for Oracle like databases")
) -> InstanceDetail:
    client = create_client()
    req = dms_enterprise_20181101_models.GetInstanceRequest(host=host, port=port)
    if sid: req.sid = sid
    if mcp.state.real_login_uid:
        req.real_login_user_uid = mcp.state.real_login_uid
    try:
        resp = client.get_instance(req)
        instance_data = resp.body.to_map().get('Instance', {}) if resp and resp.body else {}
        instance_data['InstanceResourceId'] = instance_data.pop('EcsInstanceId', None)

        return InstanceDetail(**instance_data)
    except Exception as e:
        logger.error(f"Error in get_instance: {e}")
        raise


async def list_instance(
        search_key: Optional[str] = Field(default=None,
                                          description="Optional search key (e.g., instance host, instance alias, etc.)"),
        db_type: Optional[str] = Field(default=None,
                                       description="Optional instanceType, or called dbType (e.g., mysql, polardb, oracle, "
                                                   "postgresql, sqlserver, polardb-pg, etc.)"),
        env_type: Optional[str] = Field(default=None,
                                        description="Optional instance environment type (e.g., product, dev, test, etc. )")
) -> List[InstanceDetail]:
    client = create_client()
    req = dms_enterprise_20181101_models.ListInstancesRequest()
    if search_key:
        req.search_key = search_key
    if db_type:
        req.db_type = db_type
    if env_type:
        req.env_type = env_type
    if mcp.state.real_login_uid:
        req.real_login_user_uid = mcp.state.real_login_uid
    try:
        resp = client.list_instances(req)

        instance_data = resp.body.to_map()
        if "InstanceList" not in instance_data:
            return []
        instance_list = instance_data["InstanceList"]
        # 检查是否有 Instance 键
        if "Instance" not in instance_list:
            return []
        instances = instance_list["Instance"]
        # 检查是否为空
        if not isinstance(instances, list) or not instances:
            return []

        processed_instances = [
            {**item, 'InstanceResourceId': item.pop('EcsInstanceId', None)}
            for item in instances
        ]
        return [InstanceDetail(**item) for item in processed_instances]
    except Exception as e:
        logger.error(f"Error in list_instance: {e}")
        raise


async def search_database(
        search_key: str = Field(description="database name to search for"),
        page_number: int = Field(default=1, description="Page number (starting from 1)"),
        page_size: int = Field(default=200, description="Results per page (max 1000)")
) -> List[DatabaseInfo]:
    client = create_client()
    req = dms_enterprise_20181101_models.SearchDatabaseRequest(search_key=search_key, page_number=page_number,
                                                               page_size=page_size)
    if mcp.state.real_login_uid:
        req.real_login_user_uid = mcp.state.real_login_uid
    try:
        resp = client.search_database(req)
        if not resp or not resp.body: return []
        db_list_data = resp.body.to_map().get('SearchDatabaseList', {}).get('SearchDatabase', [])
        result = []
        for db in db_list_data:
            db_info_map = {"DatabaseId": db.get("DatabaseId"), "Host": db.get("Host"), "Port": db.get("Port"),
                           "DbType": db.get("DbType")}
            db_info_map["SchemaName"] = f'{db.get("CatalogName", "")}.{db.get("SchemaName", "")}' if db.get(
                "CatalogName") != 'def' else db.get("SchemaName", "")
            result.append(DatabaseInfo(**db_info_map))
        return result
    except Exception as e:
        logger.error(f"Error in search_database: {e}")
        raise


async def get_database(
        host: str = Field(description="Hostname or IP of the database instance"),
        port: str = Field(description="Connection port number"),
        schema_name: str = Field(description="Name of the database schema"),
        sid: Optional[str] = Field(default=None, description="Required for Oracle like databases")
) -> DatabaseDetail:
    client = create_client()
    req = dms_enterprise_20181101_models.GetDatabaseRequest(host=host, port=port, schema_name=schema_name)

    if sid:
        req.sid = sid
    if mcp.state.real_login_uid:
        req.real_login_user_uid = mcp.state.real_login_uid

    try:
        resp = client.get_database(req)
        db_data = resp.body.to_map().get('Database', {}) if resp and resp.body else {}
        return DatabaseDetail(**db_data)
    except Exception as e:
        logger.error(f"Error in get_database: {e}")
        raise


async def list_tables(  # Renamed from listTable to follow convention
        database_id: str = Field(description="DMS databaseId"),
        search_name: Optional[str] = Field(default=None, description="Optional: Search keyword for table names"),
        page_number: int = Field(default=1, description="Pagination page number"),
        page_size: int = Field(default=200, description="Results per page (max 200)")
) -> Dict[str, Any]:
    if not search_name:
        search_name = "%"
    client = create_client()
    req = dms_enterprise_20181101_models.ListTablesRequest(database_id=database_id, search_name=search_name,
                                                           page_number=page_number, page_size=page_size,
                                                           return_guid=True)
    if mcp.state.real_login_uid:
        req.real_login_user_uid = mcp.state.real_login_uid
    try:
        resp = client.list_tables(req)
        return resp.body.to_map() if resp and resp.body else {}
    except Exception as e:
        logger.error(f"Error in list_tables: {e}")
        raise


async def get_meta_table_detail_info(
        table_guid: str = Field(
            description="Unique table identifier (format: dmsTableId.schemaName.tableName),Example: IDB_1567890.mySchema.myTable")
) -> TableDetail:
    client = create_client()
    req = dms_enterprise_20181101_models.GetMetaTableDetailInfoRequest(table_guid=table_guid)
    if mcp.state.real_login_uid:
        req.real_login_user_uid = mcp.state.real_login_uid
    try:
        resp = client.get_meta_table_detail_info(req)
        detail_info = resp.body.to_map().get('DetailInfo', {}) if resp and resp.body else {}
        return TableDetail(**detail_info)
    except Exception as e:
        logger.error(f"Error in get_meta_table_detail_info: {e}")
        raise


def _format_as_markdown_table(column_names: List[str], rows: List[Dict[str, Any]]) -> str:
    if not column_names: return ""
    header = "| " + " | ".join(column_names) + " |"
    separator = "| " + " | ".join(["---"] * len(column_names)) + " |"
    table_rows_str = [header, separator]
    for row_data in rows:
        row_values = [str(row_data.get(col, "")) for col in column_names]
        table_rows_str.append("| " + " | ".join(row_values) + " |")
    return "\n".join(table_rows_str)


async def execute_script(
        database_id: str = Field(description="DMS databaseId"),
        script: str = Field(description="SQL script to execute"),
        logic: bool = Field(default=False, description="Whether to use logical execution mode")
) -> ExecuteScriptResult:  # Return the object, __str__ will be used by wrapper if needed
    client = create_client()
    req = dms_enterprise_20181101_models.ExecuteScriptRequest(db_id=database_id, script=script, logic=logic)
    if mcp.state.real_login_uid:
        req.real_login_user_uid = mcp.state.real_login_uid
    try:
        resp = client.execute_script(req)
        if not resp or not resp.body:
            return ExecuteScriptResult(RequestId="", Results=[], Success=False)
        data = resp.body.to_map()
        processed_results = []
        if data.get('Success') and data.get('Results'):
            for res_item in data.get('Results', []):
                if res_item.get('Success'):
                    column_names = res_item.get('ColumnNames', [])
                    rows_data = res_item.get('Rows', [])
                    markdown_table = _format_as_markdown_table(column_names, rows_data)
                    processed_results.append(
                        ResultSet(ColumnNames=column_names, RowCount=res_item.get('RowCount', 0), Rows=rows_data,
                                  MarkdownTable=markdown_table, Success=True, Message=''))
                else:
                    processed_results.append(
                        ResultSet(ColumnNames=[], RowCount=0, Rows=[], MarkdownTable=None, Success=False,
                                  Message=res_item.get('Message')))
        return ExecuteScriptResult(RequestId=data.get('RequestId', ""), Results=processed_results,
                                   Success=data.get('Success', False))
    except Exception as e:
        logger.error(f"Error in execute_script: {e}")
        if "The instance is not in secure hosting mode" in str(e):
            return "当前实例尚未开启安全托管功能。您可以通过DMS控制台免费开启「安全托管模式」。请注意，该操作需要管理员或DBA身份权限。"


async def create_data_change_order(
        database_id: str = Field(description="DMS databaseId"),
        script: str = Field(description="SQL script to execute"),
        logic: bool = Field(default=False, description="Whether to use logical execution mode")
) -> Dict[str, Any]:
    client = create_client()
    req = dms_enterprise_20181101_models.CreateDataCorrectOrderRequest()
    req.comment = "Data correct order submitted by MCP"

    param = dms_enterprise_20181101_models.CreateDataCorrectOrderRequestParam()
    param.estimate_affect_rows = 1
    param.sql_type = "TEXT"
    param.exec_sql = script
    param.classify = "MCP"

    db_list = dms_enterprise_20181101_models.CreateDataCorrectOrderRequestParamDbItemList()
    db_list.db_id = database_id
    db_list.logic = logic

    db_items = [db_list]
    param.db_item_list = db_items

    req.param = param
    try:
        resp = client.create_data_correct_order(req)
        return resp.body.to_map()
    except Exception as e:
        logger.error(f"Error in create_data_change_order: {e}")
        raise


async def get_order_base_info(
        order_id: str = Field(description="DMS order ID")
) -> Dict[str, Any]:
    client = create_client()
    req = dms_enterprise_20181101_models.GetOrderBaseInfoRequest()
    req.order_id = order_id
    try:
        resp = client.get_order_base_info(req)
        return resp.body.to_map()
    except Exception as e:
        logger.error(f"Error in get_order_base_info: {e}")
        raise


async def submit_order_approval(
        order_id: str = Field(description="DMS order ID")
) -> Dict[str, Any]:
    client = create_client()
    req = dms_enterprise_20181101_models.SubmitOrderApprovalRequest()
    req.order_id = order_id
    try:
        resp = client.submit_order_approval(req)
        return resp.body.to_map()
    except Exception as e:
        logger.error(f"Error in submit_order_approval: {e}")
        raise


async def nl2sql(
        database_id: str = Field(description="DMS databaseId"),
        question: str = Field(description="Natural language question"),
        knowledge: Optional[str] = Field(default=None, description="Optional: additional context"),
        model: Optional[str] = Field(default=None,
                                     description="Optional: if a specific model is desired, it can be specified here")
) -> SqlResult:
    client = create_client()
    req = dms_enterprise_20181101_models.GenerateSqlFromNLRequest(db_id=database_id, question=question)
    if knowledge:
        req.knowledge = knowledge
    if mcp.state.real_login_uid:
        req.real_login_user_uid = mcp.state.real_login_uid
    if model:
        req.model = model
    try:
        resp = client.generate_sql_from_nl(req)
        if not resp or not resp.body: return SqlResult(sql=None)
        data = resp.body.to_map()
        sql_content = data.get('Data', {}).get('Sql') if data else None
        return SqlResult(sql=sql_content)
    except Exception as e:
        logger.error(f"Error in nl2sql_explicit_db: {e}")
        raise


async def answer_sql_syntax(
        database_id: str = Field(description="DMS databaseId"),
        question: str = Field(description="Natural language question"),
        model: Optional[str] = Field(default=None,
                                     description="Optional: if a specific model is desired, it can be specified here")
) -> Dict[str, Any]:
    client = create_client()
    req = dms_enterprise_20181101_models.AnswerSqlSyntaxByMetaAgentRequest(db_id=database_id, query=question)
    # if mcp.state.real_login_uid:
    #     req.real_login_user_uid = mcp.state.real_login_uid
    if model:
        req.model = model
    try:
        resp = client.answer_sql_syntax_by_meta_agent(req)
        if not resp or not resp.body:
            return None
        data = resp.body.to_map()
        return data
    except Exception as e:
        logger.error(f"Error in ask_sql_syntax: {e}")
        raise


async def fix_sql_statement(
        database_id: str = Field(description="DMS databaseId"),
        question: Optional[str] = Field(default=None, description="Natural language question"),
        sql: str = Field(description="The SQL that caused an error"),
        error: str = Field(description="SQL error message"),
        model: Optional[str] = Field(default=None,
                                     description="Optional: if a specific model is desired, it can be specified here")
) -> Dict[str, Any]:
    client = create_client()
    req = dms_enterprise_20181101_models.FixSqlByMetaAgentRequest(db_id=database_id, query=question, sql=sql,
                                                                  error=error)
    # if mcp.state.real_login_uid:
    #     req.real_login_user_uid = mcp.state.real_login_uid
    if model:
        req.model = model
    try:
        resp = client.fix_sql_by_meta_agent(req)
        if not resp or not resp.body:
            return None
        data = resp.body.to_map()
        return data
    except Exception as e:
        logger.error(f"Error in fix_sql_statement: {e}")
        raise


async def optimize_sql(
        database_id: str = Field(description="DMS databaseId"),
        question: Optional[str] = Field(default=None, description="Natural language question"),
        sql: str = Field(description="SQL statement"),
        model: Optional[str] = Field(default=None,
                                     description="Optional: if a specific model is desired, it can be specified here")
) -> Any:
    client = create_client()
    req = dms_enterprise_20181101_models.OptimizeSqlByMetaAgentRequest(db_id=database_id, query=question, sql=sql)
    # if mcp.state.real_login_uid:
    #     req.real_login_user_uid = mcp.state.real_login_uid
    if model:
        req.model = model
    try:
        resp = client.optimize_sql_by_meta_agent(req)
        if not resp or not resp.body:
            return None
        data = resp.body.to_map()
        return data
    except Exception as e:
        logger.error(f"Error in optimize_sql: {e}")
        raise


# --- ToolRegistry Class ---
class ToolRegistry:
    def __init__(self, mcp: FastMCP):
        self.mcp = mcp
        self.default_database_id: Optional[str] = getattr(self.mcp.state, 'default_database_id', None)

    def register_tools(self) -> FastMCP:
        if self.default_database_id:
            logger.info(f"DATABASE_ID is set ('{self.default_database_id}'). Registering configured toolset.")
            self._register_configured_db_toolset()
        else:
            logger.info("DATABASE_ID not set. Registering full toolset.")
            self._register_full_toolset()
        return self.mcp

    def _register_configured_db_toolset(self):
        @self.mcp.tool(name="listTables",
                       description="Lists tables in the database. Search by name is supported.",
                       annotations={"title": "List Tables (Pre-configured DB)", "readOnlyHint": True})
        async def list_tables_configured(
                search_name: Optional[str] = Field(default=None,
                                                   description="Optional: A string used as the search keyword to match table names."),
                page_number: int = Field(description="Pagination page number", default=1),
                page_size: int = Field(description="Number of results per page", default=200)
        ) -> Dict[str, Any]:
            if not search_name:
                search_name = "%"
            return await list_tables(database_id=self.default_database_id, search_name=search_name,
                                     page_number=page_number, page_size=page_size)

        self.mcp.tool(name="getTableDetailInfo",
                      description="Retrieve detailed metadata information about a specific database table including "
                                  "schema and index details. If you don't know the table_guid parameter, retrieve it using listTables.",
                      annotations={"title": "Get Table Details", "readOnlyHint": True})(get_meta_table_detail_info)

        @self.mcp.tool(name="executeScript",
                       description="Executes an SQL script against the pre-configured database.",
                       annotations={"title": "Execute SQL (Pre-configured DB)", "readOnlyHint": False,
                                    "destructiveHint": True})
        async def execute_script_configured(
                script: str = Field(description="SQL script to execute")
        ) -> str:
            result_obj = await execute_script(database_id=self.default_database_id, script=script, logic=False)
            return str(result_obj)

        @self.mcp.tool(name="createDataChangeOrder",
                       description="Execute SQL changes through a data change order, and a corresponding order ID will be returned. "
                                   "Prefer using the executeScript tool for SQL execution; "
                                   "only use this tool when explicitly instructed to perform the operation via a order.",
                       annotations={"title": "在DMS中创建数据变更工单", "readOnlyHint": False, "destructiveHint": True})
        async def create_data_change_order_configured(
                script: str = Field(description="SQL script to execute")
        ) -> str:
            result_obj = await create_data_change_order(database_id=self.default_database_id, script=script,
                                                        logic=False)
            return str(result_obj)

        self.mcp.tool(name="getOrderInfo", description="Retrieve order information from DMS using the order ID.",
                      annotations={"title": "获取DMS工单详情", "readOnlyHint": True})(get_order_base_info)

        self.mcp.tool(name="submitOrderApproval",
                      description="Submit the order for approval in DMS using the order ID.",
                      annotations={"title": "提交工单审批", "readOnlyHint": False})(submit_order_approval)

        @self.mcp.tool(name="askDatabase",
                       description="Ask a question in natural language to the pre-configured database and get results directly.",
                       annotations={"title": "Ask Pre-configured Database", "readOnlyHint": True})
        async def ask_database_configured(
                question: str = Field(
                    description="Your question in natural language about the pre-configured database."),
                knowledge: Optional[str] = Field(default=None,
                                                 description="Optional: additional context to help formulate the SQL query."),
                model: Optional[str] = Field(default=None,
                                             description="Optional: if a specific model is desired, it can be specified here")
        ) -> AskDatabaseResult:
            sql_result_obj = await nl2sql(database_id=self.default_database_id, question=question,
                                          knowledge=knowledge, model=model)
            generated_sql = ""
            if not sql_result_obj or not sql_result_obj.sql:
                logger.warning(f"Failed to generate SQL for question: {question} on preconfigured DB.")
                return AskDatabaseResult(executed_sql=generated_sql,
                                         execution_result="Error: Could not generate an SQL query from your question.")

            generated_sql = sql_result_obj.sql
            logger.info(f"Generated SQL for pre-configured DB: {generated_sql}")
            try:
                execution_result_obj = await execute_script(database_id=self.default_database_id, script=generated_sql,
                                                            logic=False)
                return AskDatabaseResult(executed_sql=generated_sql, execution_result=str(execution_result_obj))
            except Exception as e:
                logger.error(f"Error executing SQL for pre-configured DB: {e}")
                return AskDatabaseResult(executed_sql=generated_sql,
                                         execution_result=f"Error: An issue occurred while executing the query: {str(e)}")

        @self.mcp.tool(name="answerSqlSyntax",
                       description="Answer syntax-related questions for the corresponding database engine ",
                       annotations={"title": "SQL语法回答", "readOnlyHint": True, "destructiveHint": False})
        async def answer_sql_syntax_configured(
                question: str = Field(description="Natural language question"),
                model: Optional[str] = Field(default=None,
                                             description="Optional: if a specific model is desired, it can be specified here")
        ) -> Dict[str, Any]:
            result_obj = await answer_sql_syntax(database_id=self.default_database_id, question=question,
                                                 model=model)
            return result_obj

        @self.mcp.tool(name="fixSql",
                       description="Analyze and fix the SQL error based on the provided SQL statement and error message.",
                       annotations={"title": "SQL修复", "readOnlyHint": True, "destructiveHint": False})
        async def fix_sql_configured(
                question: Optional[str] = Field(default=None, description="Natural language question"),
                sql: str = Field(description="The SQL that caused an error"),
                error: str = Field(description="SQL error message"),
                model: Optional[str] = Field(default=None,
                                             description="Optional: if a specific model is desired, it can be specified here")
        ) -> Dict[str, Any]:
            result_obj = await fix_sql_statement(database_id=self.default_database_id, question=question, sql=sql,
                                                 error=error, model=model)
            return result_obj

        @self.mcp.tool(name="optimizeSql",
                       description="Analyze and optimize SQL performance based on the provided SQL statement",
                       annotations={"title": "SQL优化", "readOnlyHint": True, "destructiveHint": False})
        async def optimize_sql_configured(
                question: Optional[str] = Field(default=None, description="Natural language question"),
                sql: str = Field(description="SQL statement"),
                model: Optional[str] = Field(default=None,
                                             description="Optional: if a specific model is desired, it can be specified here")
        ) -> Any:
            result_obj = await optimize_sql(database_id=self.default_database_id, question=question, sql=sql,
                                            model=model)
            return result_obj

    def _register_full_toolset(self):
        self.mcp.tool(name="addInstance",
                      description="Add an instance to DMS. The username and password are required. "
                                  "Only Aliyun instances are supported. "
                                  "Either instance_resource_id or host and port must be provided. "
                                  "The region is optional, but it's recommended to include it."
                                  "If the instance already exists, it will return the existing instance information.",
                      annotations={"title": "添加或获取DMS实例", "readOnlyHint": False, "destructiveHint": False})(
            add_instance)
        self.mcp.tool(name="listInstances", description="Search for instances from DMS.",
                      annotations={"title": "搜索DMS实例列表", "readOnlyHint": True})(list_instance)
        self.mcp.tool(name="getInstance",
                      description="Retrieve detailed instance information from DMS using the host and port.",
                      annotations={"title": "获取DMS实例详情", "readOnlyHint": True})(get_instance)
        self.mcp.tool(name="searchDatabase", description="Search databases in DMS by name.",
                      annotations={"title": "搜索DMS数据库", "readOnlyHint": True})(search_database)
        self.mcp.tool(name="getDatabase",
                      description="Obtain detailed information about a specific database in DMS when the host and port are provided.",
                      annotations={"title": "获取DMS数据库详情", "readOnlyHint": True})(get_database)
        self.mcp.tool(name="listTables",
                      description=f"Search for tables by databaseId and (optional) table name. "
                                  f"{DATABASE_ID_DESCRIPTION}",
                      annotations={"title": "列出DMS表", "readOnlyHint": True})(list_tables)
        self.mcp.tool(name="getTableDetailInfo",
                      description="Retrieve detailed metadata information about a specific database table including "
                                  "schema and index details. If you don't know the table_guid parameter, retrieve it using listTables.",
                      annotations={"title": "获取DMS表详细信息", "readOnlyHint": True})(get_meta_table_detail_info)

        @self.mcp.tool(name="executeScript",
                       description=f"Execute SQL script against a database in DMS and return structured results."
                                   f"{DATABASE_ID_DESCRIPTION}",
                       annotations={"title": "在DMS中执行SQL脚本", "readOnlyHint": False, "destructiveHint": True})
        async def execute_script_full_wrapper(
                database_id: str = Field(description="Required DMS databaseId. Obtained via getDatabase tool"),
                script: str = Field(description="SQL script to execute"),
                logic: bool = Field(description="Whether to use logical execution mode", default=False)
        ) -> str:  # Return string representation
            result_obj = await execute_script(database_id=database_id, script=script, logic=logic)
            return str(result_obj)

        @self.mcp.tool(name="createDataChangeOrder",
                       description=f"Execute SQL changes through a data change order, and a corresponding order ID will be returned. "
                                   f"Prefer using the executeScript tool for SQL execution;"
                                   f"only use this tool when explicitly instructed to perform the operation via a order."
                                   f"{DATABASE_ID_DESCRIPTION}",
                       annotations={"title": "在DMS中创建数据变更工单", "readOnlyHint": False, "destructiveHint": True})
        async def create_data_change_order_wrapper(
                database_id: str = Field(description="Required DMS databaseId. Obtained via getDatabase tool"),
                script: str = Field(description="SQL script to execute"),
                logic: bool = Field(description="Whether to use logical execution mode", default=False)
        ) -> str:  # Return string representation
            result_obj = await create_data_change_order(database_id=database_id, script=script, logic=logic)
            return str(result_obj)

        self.mcp.tool(name="getOrderInfo", description="Retrieve order information from DMS using the order ID.",
                      annotations={"title": "获取DMS工单详情", "readOnlyHint": True})(get_order_base_info)

        self.mcp.tool(name="submitOrderApproval",
                      description="Submit the order for approval in DMS using the order ID.",
                      annotations={"title": "提交工单审批", "readOnlyHint": False})(submit_order_approval)

        self.mcp.tool(name="generateSql", description="Generate SELECT-type SQL queries from natural language input.",
                      annotations={"title": "自然语言转SQL (DMS)", "readOnlyHint": True})(nl2sql)
        self.mcp.tool(name="fixSql", description=f"Analyze and fix the SQL error based on the provided "
                                                 f"SQL statement, error message, and database ID."
                                                 f"{DATABASE_ID_DESCRIPTION}",
                      annotations={"title": "SQL修复", "readOnlyHint": True})(fix_sql_statement)
        self.mcp.tool(name="answerSqlSyntax", description=f"Answer syntax-related questions "
                                                          f"for the corresponding database engine "
                                                          f"based on the database ID."
                                                          f"{DATABASE_ID_DESCRIPTION}",
                      annotations={"title": "SQL语法回答", "readOnlyHint": True})(answer_sql_syntax)
        self.mcp.tool(name="optimizeSql", description=f"Analyze and optimize SQL performance "
                                                      f"based on the provided SQL statement and database ID"
                                                      f"{DATABASE_ID_DESCRIPTION}",
                      annotations={"title": "SQL优化", "readOnlyHint": True})(optimize_sql)


# --- Lifespan Function ---
@asynccontextmanager
async def lifespan(app: FastMCP) -> AsyncGenerator[None, None]:
    logger.info("Initializing DMS MCP Server via lifespan")

    # Ensure app.state exists
    if not hasattr(app, 'state') or app.state is None:
        class AppState: pass

        app.state = AppState()

    # Initialize realLoginUid
    app.state.real_login_uid = None
    uid = os.getenv("UID")
    if uid:
        app.state.real_login_uid = uid
        logger.info(f"RealLoginUid environment variable found: {uid}")

    # Initialize default_database_id
    app.state.default_database_id = None

    dms_connection_string = os.getenv("CONNECTION_STRING")
    if dms_connection_string:
        logger.info(f"CONNECTION_STRING environment variable found: {dms_connection_string}")
        db_host, db_port, db_name_path, catalog_name = None, None, None, None
        try:
            # Expected formats:
            # 1. catalog@host:port:schema  (PG, full)
            # 2. database@host:port      (MySQL, database is schema-like)
            # 3. host:port:schema          (No catalog, with schema)
            # 4. host:port                 (No catalog, no schema)

            parts = dms_connection_string.split('@')

            potential_catalog_or_db_name = None
            main_part = ""

            if len(parts) > 1:  # Contains '@'
                potential_catalog_or_db_name = parts[0]
                main_part = parts[1]
            else:  # No '@'
                main_part = parts[0]

            main_part_components = main_part.split(':')

            if len(main_part_components) == 3:  # host:port:schema
                db_host = main_part_components[0]
                db_port = main_part_components[1]
                db_name_path = main_part_components[2]  # This is schema
                if potential_catalog_or_db_name:  # Format 1: catalog@host:port:schema
                    catalog_name = potential_catalog_or_db_name
                # else Format 3: host:port:schema (catalog_name remains None)

            elif len(main_part_components) == 2:  # host:port
                db_host = main_part_components[0]
                db_port = main_part_components[1]
                if potential_catalog_or_db_name:  # Format 2: database@host:port
                    # Here, potential_catalog_or_db_name is the database/schema
                    db_name_path = potential_catalog_or_db_name
                    # For MySQL-like, catalog is not explicit in this way, so catalog_name remains None or is not used as such.
                # else Format 4: host:port (db_name_path and catalog_name remain None)

            else:
                raise ValueError(
                    f"Invalid format for host:port or host:port:schema part: '{main_part}' in CONNECTION_STRING '{dms_connection_string}'.")

            if not (db_host and db_port):  # This check might be redundant if ValueError above catches it.
                logger.error(
                    f"CONNECTION_STRING '{dms_connection_string}' is incomplete. Missing host or port. Expected formats: catalog@host:port:schema, database@host:port, host:port:schema, host:port")
            else:
                logger.info(f"Verifying instance from CONNECTION_STRING: {db_host}:{db_port}")
                try:
                    instance_details = await get_instance(host=db_host, port=str(db_port), sid=None)
                    if not instance_details or not hasattr(instance_details,
                                                           'InstanceId') or not instance_details.InstanceId:
                        logger.warning(
                            f"Instance {db_host}:{db_port} not found or no valid InstanceId returned by get_instance. Cannot use this CONNECTION_STRING.")
                    else:
                        logger.info(f"Instance {db_host}:{db_port} verified. InstanceId: {instance_details.InstanceId}")

                        if db_name_path or catalog_name:  # We need either a schema or a catalog to search
                            search_term_for_db = catalog_name if catalog_name else db_name_path
                            sid_name = None
                            if catalog_name and db_name_path:
                                sid_name = db_name_path
                                logger.info(
                                    f"Searching for database with catalog '{catalog_name}' and schema '{db_name_path}' associated with instance {db_host}:{db_port}")
                            elif db_name_path:
                                logger.info(
                                    f"Searching for database schema '{db_name_path}' associated with instance {db_host}:{db_port}")
                            elif catalog_name:
                                logger.info(
                                    f"Searching for database catalog '{catalog_name}' associated with instance {db_host}:{db_port}")

                            database = await get_database(host=db_host,
                                                          port=db_port, schema_name=search_term_for_db, sid=sid_name)
                            if not database or database.DatabaseId is None:
                                logger.warning(
                                    f"No database found for {search_term_for_db} at {db_host}:{db_port} after processing CONNECTION_STRING.")
                                database = await get_database(host=db_host, port=db_port,
                                                              schema_name=search_term_for_db, sid=None)
                            found_db_id = None
                            if database:
                                found_db_id = database.DatabaseId

                            if found_db_id:
                                app.state.default_database_id = found_db_id
                                logger.info(
                                    f"Successfully configured default_database_id to {found_db_id} using CONNECTION_STRING.")
                            else:
                                current_search_criteria = f"catalog '{catalog_name}', schema '{db_name_path}'" if catalog_name and db_name_path else f"schema '{db_name_path}'" if db_name_path else f"catalog '{catalog_name}'"
                                logger.warning(
                                    f"Could not find a matching database for {current_search_criteria} at {db_host}:{db_port} after processing CONNECTION_STRING.")
                        else:
                            logger.info(
                                f"Instance {db_host}:{db_port} verified, but no catalog or schema provided in CONNECTION_STRING. No default database_id will be set from this DSN.")

                except Exception as instance_e:
                    logger.error(
                        f"Error during instance verification or database search for CONNECTION_STRING '{dms_connection_string}': {instance_e}")

        except ValueError as ve:
            logger.error(
                f"Invalid CONNECTION_STRING format '{dms_connection_string}': {ve}. Expected formats: catalog@host:port:schema, database@host:port, host:port:schema, or host:port")
        except Exception as e:
            logger.error(f"Error parsing CONNECTION_STRING '{dms_connection_string}': {e}")
    else:
        logger.info("CONNECTION_STRING environment variable not found.")

    if app.state.default_database_id:
        logger.info(f"Final default_database_id to be used (from CONNECTION_STRING): {app.state.default_database_id}")
    else:
        logger.info("No default database ID configured from CONNECTION_STRING. Full toolset will be registered.")

    registry = ToolRegistry(mcp=app)
    registry.register_tools()

    yield

    logger.info("Shutting down DMS MCP Server via lifespan")
    if hasattr(app.state, 'default_database_id'):
        delattr(app.state, 'default_database_id')
    if hasattr(app.state, 'real_login_uid'):
        delattr(app.state, 'real_login_uid')


# --- FastMCP Instance Creation & Server Run ---
mcp = FastMCP(
    "DatabaseManagementAssistant",
    lifespan=lifespan,
    instructions="Database Management Assistant (DMS) is a toolkit designed to assist users in managing and "
                 "interacting with databases."
)


def run_server():
    log_level_str = os.getenv("LOG_LEVEL", "INFO").upper()
    logging.basicConfig(level=log_level_str, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    logger.info(f"Starting DMS MCP server with log level {log_level_str}")
    mcp.run(transport=os.getenv('SERVER_TRANSPORT', 'stdio'))


if __name__ == "__main__":
    run_server()
