import os
import json
import time
from types import SimpleNamespace
from urllib.parse import quote_plus
from urllib import error, request
from dotenv import load_dotenv
from sqlalchemy import create_engine
import pandas as pd
from mcp.server.fastmcp import FastMCP

# Load environment variables
load_dotenv()

# Initialize FastMCP server
mcp = FastMCP("mysql-mcp")
_feishu_tat_cache = {"token": "", "expire_at": 0.0, "app_id": "", "app_secret": ""}

def get_connect_args(configs):
    """
    Get database connection arguments, supporting SSL and other generic configurations
    """
    connect_args = {}
    
    # Check if SSL is required
    use_ssl = getattr(configs, 'use_ssl', False)
    
    # TiDB defaults to using SSL if specified or if db_type is tidb (though here we simplify)
    if use_ssl:
        ssl_ca = getattr(configs, 'ssl_ca', '')
        if ssl_ca:
            connect_args = {
                'ssl': {
                    'ca': ssl_ca,
                    'check_hostname': False,
                    'verify_identity': False
                }
            }
    
    return connect_args

def get_mysql_config():
    """Create a config object from environment variables."""
    return SimpleNamespace(
        host=os.getenv("MYSQL_HOST", "localhost"),
        port=int(os.getenv("MYSQL_PORT", 3306)),
        user=os.getenv("MYSQL_USER", "root"),
        password=os.getenv("MYSQL_PASSWORD", ""),
        database=os.getenv("MYSQL_DATABASE", ""),
        use_ssl=os.getenv("MYSQL_USE_SSL", "false").lower() == "true",
        ssl_ca=os.getenv("MYSQL_SSL_CA", ""),
    )

def get_feishu_config():
    return SimpleNamespace(
        mcp_url=os.getenv("FEISHU_MCP_URL", "https://mcp.feishu.cn/mcp"),
        app_id=os.getenv("FEISHU_MCP_APP_ID", ""),
        app_secret=os.getenv("FEISHU_MCP_APP_SECRET", ""),
        allowed_tools=os.getenv("FEISHU_MCP_ALLOWED_TOOLS", ""),
    )

def get_feishu_tat(configs):
    app_id = configs.app_id.strip()
    app_secret = configs.app_secret.strip()
    if not app_id or not app_secret:
        raise ValueError("Missing Feishu app credential. Set FEISHU_MCP_APP_ID and FEISHU_MCP_APP_SECRET.")

    now = time.time()
    if (
        _feishu_tat_cache["token"]
        and _feishu_tat_cache["expire_at"] > now + 60
        and _feishu_tat_cache["app_id"] == app_id
        and _feishu_tat_cache["app_secret"] == app_secret
    ):
        return _feishu_tat_cache["token"]

    token_url = "https://open.feishu.cn/open-apis/auth/v3/tenant_access_token/internal"
    body = json.dumps({"app_id": app_id, "app_secret": app_secret}).encode("utf-8")
    req = request.Request(token_url, data=body, headers={"Content-Type": "application/json"}, method="POST")

    try:
        with request.urlopen(req, timeout=30) as resp:
            response_text = resp.read().decode("utf-8")
    except error.HTTPError as e:
        detail = e.read().decode("utf-8", errors="ignore")
        raise RuntimeError(detail or str(e)) from e
    except Exception as e:
        raise RuntimeError(str(e)) from e

    try:
        data = json.loads(response_text)
    except json.JSONDecodeError as e:
        raise RuntimeError(response_text) from e

    if data.get("code") != 0:
        raise RuntimeError(data.get("msg") or data.get("message") or "Failed to get tenant access token.")

    token = data.get("tenant_access_token", "")
    expire = int(data.get("expire", 0))
    if not token:
        raise RuntimeError("Feishu auth response missing tenant_access_token.")

    _feishu_tat_cache["token"] = token
    _feishu_tat_cache["expire_at"] = now + max(expire, 0)
    _feishu_tat_cache["app_id"] = app_id
    _feishu_tat_cache["app_secret"] = app_secret
    return token

def build_feishu_headers(configs, allowed_tools: str = ""):
    headers = {"Content-Type": "application/json"}
    headers["X-Lark-MCP-TAT"] = get_feishu_tat(configs)

    resolved_allowed_tools = allowed_tools.strip() if allowed_tools else configs.allowed_tools.strip()
    if resolved_allowed_tools:
        tools = [tool.strip() for tool in resolved_allowed_tools.split(",") if tool.strip()]
        if tools:
            headers["X-Lark-MCP-Allowed-Tools"] = ",".join(tools)
    return headers

def call_feishu_rpc(method: str, params: dict | None, configs, allowed_tools: str = "") -> dict:
    payload = {"jsonrpc": "2.0", "id": 1, "method": method}
    if params is not None:
        payload["params"] = params
    body = json.dumps(payload).encode("utf-8")
    headers = build_feishu_headers(configs, allowed_tools=allowed_tools)
    req = request.Request(configs.mcp_url, data=body, headers=headers, method="POST")

    try:
        with request.urlopen(req, timeout=30) as resp:
            response_text = resp.read().decode("utf-8")
    except error.HTTPError as e:
        detail = e.read().decode("utf-8", errors="ignore")
        return {"jsonrpc": "2.0", "id": 1, "error": {"code": e.code, "message": detail or str(e)}}
    except Exception as e:
        return {"jsonrpc": "2.0", "id": 1, "error": {"code": -1, "message": str(e)}}

    try:
        return json.loads(response_text)
    except json.JSONDecodeError:
        return {"jsonrpc": "2.0", "id": 1, "error": {"code": -2, "message": response_text}}

def execute_query(sql, configs):
    """
    Execute SQL query and return DataFrame
    """
    try:
        cf_port = configs.port
    except:
        cf_port = 3306
        
    # URL-encode the password to handle special characters
    password_encoded = quote_plus(configs.password)
    
    # Get connection arguments
    connect_args = get_connect_args(configs)
    
    # Create engine
    connection_string = f'mysql+pymysql://{configs.user}:{password_encoded}@{configs.host}:{cf_port}/{configs.database}'
    engine = create_engine(connection_string, connect_args=connect_args)
    
    # Connect and read
    conn = engine.raw_connection()
    try:
        df = pd.read_sql(sql, conn)
        return df
    finally:
        conn.close()

@mcp.tool()
def read_sql(query_sql: str) -> str:
    """
    Execute a read-only SQL query on the configured MySQL database.
    Only SELECT and SHOW statements are allowed for safety.
    
    Args:
        query_sql: The SQL query to execute (e.g., "SELECT * FROM users LIMIT 10")
        
    Returns:
        A JSON string representation of the query result.
    """
    # Basic security check for read-only
    normalized_sql = query_sql.strip().upper()
    if not (normalized_sql.startswith("SELECT") or normalized_sql.startswith("SHOW")):
        return "Error: Only SELECT and SHOW statements are allowed in this read-only tool."

    try:
        config = get_mysql_config()
        
        df = execute_query(query_sql, config)
        
        if df is None or df.empty:
            return "[]"
            
        # Convert DataFrame to JSON string
        # orient='records' creates a list of dicts: [{'col1': val1, 'col2': val2}, ...]
        result_json = df.to_json(orient="records", date_format="iso", force_ascii=False)
        return result_json
    except Exception as e:
        return f"Error executing query: {str(e)}"

@mcp.tool()
def feishu_list_tools(allowed_tools: str = "") -> str:
    try:
        config = get_feishu_config()
        init_result = call_feishu_rpc("initialize", None, config, allowed_tools=allowed_tools)
        if init_result.get("error"):
            return json.dumps(init_result, ensure_ascii=False)

        result = call_feishu_rpc("tools/list", None, config, allowed_tools=allowed_tools)
        return json.dumps(result, ensure_ascii=False)
    except Exception as e:
        return f"Error listing Feishu tools: {str(e)}"

@mcp.tool()
def feishu_call_tool(tool_name: str, arguments_json: str = "{}", allowed_tools: str = "") -> str:
    if not tool_name or not tool_name.strip():
        return "Error: tool_name is required."

    try:
        arguments = json.loads(arguments_json)
    except json.JSONDecodeError:
        return "Error: arguments_json must be valid JSON."

    if not isinstance(arguments, dict):
        return "Error: arguments_json must be a JSON object."

    try:
        config = get_feishu_config()
        init_result = call_feishu_rpc("initialize", None, config, allowed_tools=allowed_tools)
        if init_result.get("error"):
            return json.dumps(init_result, ensure_ascii=False)

        params = {"name": tool_name.strip(), "arguments": arguments}
        result = call_feishu_rpc("tools/call", params, config, allowed_tools=allowed_tools)
        return json.dumps(result, ensure_ascii=False)
    except Exception as e:
        return f"Error calling Feishu tool: {str(e)}"

def main():
    """Entry point for the package."""
    mcp.run()

if __name__ == "__main__":
    main()
