from utility import set_key

app = Flask(__name__)
app.secret_key = secrets.token_urlsafe(16)

CLIENT_ID = "40711518457*******"
REDIRECT_URI = "http://127.0.0.1:5000/oauth/callback"

DISCOVERY_URL = "https://openapi-mcp.cn-hangzhou.aliyuncs.com/.well-known/oauth-authorization-server"

def fetch_discovery_info():
    """从 discovery url 获取 Oauth 端点信息"""
    try:
        resp = requests.get(DISCOVERY_URL, timeout=5)
        if resp.status_code == 200:
            data = resp.json()
            return {
                "authorization_endpoint": data.get("authorization_endpoint"),
                "token_endpoint": data.get("token_endpoint"),
                "registration_endpoint": data.get("registration_endpoint")
            }
    except Exception as e:
        print(f"Failed to fetch discovery info: {e}")
    return {}

# 默认端点
AUTHORIZATION_ENDPOINT = "https://signin.aliyun.com/oauth2/v1/auth"
TOKEN_ENDPOINT = "https://oauth.aliyun.com/v1/token"


def generate_pkce():
    """生成 PKCE 的 code_verifier 和 code_challenge"""
    code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode().rstrip("=")
    
    # 计算 S256 code_challenge
    digest = hashlib.sha256(code_verifier.encode()).digest()
    code_challenge = base64.urlsafe_b64encode(digest).decode().rstrip("=")
    
    return code_verifier, code_challenge

@app.route("/")
def home():
    return'<a href="/login">Login with OAuth</a>'

@app.route("/login")
def login():
    registration_endpoint = ""
    # 尝试用 discovery 信息覆盖端点
    discovery = fetch_discovery_info()
    print(f"Discovery info: {discovery}")
    if discovery.get("authorization_endpoint"):
        AUTHORIZATION_ENDPOINT = discovery["authorization_endpoint"]
    if discovery.get("token_endpoint"):
        TOKEN_ENDPOINT = discovery["token_endpoint"]
    if discovery.get("registration_endpoint"):
        registration_endpoint = discovery["registration_endpoint"]
    # 注册一个 client（如果 CLIENT_ID 未设置或为占位符）
    client_id = CLIENT_ID
    if (not client_id) or client_id.endswith("*******"):
        ifnot registration_endpoint:
            return"Registration endpoint not available", 400
        # 注册 client
        reg_data = {
            "redirect_uris": [REDIRECT_URI],
            "grant_types": ["authorization_code"],
            "response_types": ["code"],
        }
        try:
            reg_resp = requests.post(registration_endpoint, json=reg_data, timeout=5)
            if reg_resp.status_code != 201:
                return f"Client registration failed: {reg_resp.text}", 400
            reg_json = reg_resp.json()
            client_id = reg_json.get("client_id")
            ifnot client_id:
                return"No client_id returned from registration", 400
            session["client_id"] = client_id
        except Exception as e:
            return f"Client registration exception: {e}", 400
    else:
        session["client_id"] = client_id

    # 生成 PKCE 参数
    code_verifier, code_challenge = generate_pkce()
    
    # 生成随机 state 防止 CSRF
    state = secrets.token_urlsafe(16)
    
    # 保存到 session
    session.update({
        "code_verifier": code_verifier,
        "state": state
    })
    
    # 构造授权请求 URL
    params = {
        "response_type": "code",
        "client_id": session["client_id"],
        "redirect_uri": REDIRECT_URI,
        "code_challenge": code_challenge,
        "code_challenge_method": "S256",
        "state": state
    }
    
    auth_url = f"{AUTHORIZATION_ENDPOINT}?{urllib.parse.urlencode(params)}"
    return redirect(auth_url)

@app.route("/oauth/callback")
def callback():
    # 检查错误响应
    if"error" in request.args:
        return f"Error: {request.args['error']}"
    
    # 验证 state
    if request.args.get("state") != session.get("state"):
        return"Invalid state parameter", 400
    
    # 获取授权码
    auth_code = request.args.get("amp;code") or request.args.get('code')  # 尝试两种可能的参数名
    ifnot auth_code:
        return"Missing authorization code", 400
    
    # 用授权码换取 token
    data = {
        "grant_type": "authorization_code",
        "code": auth_code,
        "redirect_uri": REDIRECT_URI,
        "client_id": session.get("client_id", CLIENT_ID),
        "code_verifier": session["code_verifier"]
    }
    
    response = requests.post(TOKEN_ENDPOINT, data=data)
    
    if response.status_code != 200:
        return f"Token request failed: {response.text}", 400
    
    token_info = response.json().get("access_token")
    
    # 存储到本地配置文件
    print(f"Your access_token: {token_info}")
    set_key("ALI_OPENAPI_ACCESS_TOKEN", token_info)

    # 删掉session参数
    session.pop("code_verifier", None)
    session.pop("state", None)
    
    return response.json()

if __name__ == "__main__":
    app.run(port=5000, debug=True)

print(f"Current path is {os.getcwd()}")
load_keys()

async def run_agent(message: str) -> None:

    server_params = StreamableHTTPClientParams(
        url = "https://openapi-mcp.cn-hangzhou.aliyuncs.com/accounts/1411741061209533/custom/ecs_tst_agno/id/RXPfhaBVHq7w3wkp/mcp,    # ECS
        headers = {'Authorization': f'Bearer {os.getenv("ALI_OPENAPI_ACCESS_TOKEN")}'}
    )

    async with MCPTools(server_params=server_params, transport="streamable-http", timeout_seconds=30 ) as mcp_tools:
        # Initialize the model
        model=OpenAILike(id="qwen-max", 
                         api_key=os.getenv("DASHSCOPE_API_KEY"),
                         base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")
        
        # Initialize the agent
        agent = Agent(model=model,
                      tools=[mcp_tools],
                      instructions=dedent("""\
                          你是一个阿里云云计算专家，请根据用户的问题，使用MCP服务查询阿里云的云产品信息，给出详细的解释。
                          请使用中文回答
                      """),
                      markdown=True,
                      show_tool_calls=True)

        # Run the agent
        await agent.aprint_response(message, stream=True)

# Example usage
if __name__ == "__main__":
    asyncio.run(run_agent("我在上海有哪些ECS实例?"))

def get_selected_tools_list(server_url, hearders, llm_api_key, user_question):
    from mcp.client.streamable_http import streamablehttp_client
    from mcp import ClientSession
    import asyncio
    from agno.agent import Agent, RunResponse
    from agno.models.openai.like import OpenAILike
    import json
    # Important: Just to avoid such logging error like "JSONRPCError.jsonrpc Field required ...
    import logging
    logging.disable(logging.CRITICAL)
    # 1. Get all tools via tools/list
    all_tools = None
    async def get_all_tools():
        # Connect to a streamable HTTP server
        async with streamablehttp_client(url=server_url,headers=hearders)as(read_stream, write_stream,_):
            # Create a session using the client streams
            async with ClientSession(read_stream, write_stream) as session:
                await session.initialize()
                all_tools = await session.list_tools()
                print(f"Number of tools: {len(all_tools.tools)}")
                return all_tools
    all_tools = asyncio.run(get_all_tools())
    # 2. Collect all tools brief info
    brife_tools_info = [
        {
            "name": tool.name,
            "description": tool.description,
        }
        for tool in all_tools.tools]
    # 3. Create an agent
    simple_agent = Agent(
        model=OpenAILike(
            id="qwen-max",
            api_key=llm_api_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        ),
        system_message="""
            你是一个云计算专家，严格按照请客户提供的上下文信息回答问题。
            """,
        )
   
    # 4. Run the agent to tell us which tools are needed for user questio
    prompt = (
        f"请根据提供的工具信息以及用户请求，你需要给出可能需要调用的api列表，以Json形式返回，要精简、不需要其他信息。格式上，返回值不带```、json等字符串，"
        f"工具信息如下：{json.dumps(brife_tools_info)},"
        f"现在客户提问：{user_question}"
    )
    response: RunResponse = simple_agent.run(prompt)
    print(response.content)
    tool_to_use_list = json.loads(response.content)
    print(f"Selected Tool List: {tool_to_use_list}")
    return tool_to_use_list

print(f"Current path is {os.getcwd()}")
load_keys()

async def run_agent(config):

    # Setup agent with MCP tools
    server_params = StreamableHTTPClientParams(url=config["server_url"], headers=config["headers"])
    async with MCPTools(server_params=server_params, 
                        transport="streamable-http", 
                        timeout_seconds=30, 
                        include_tools=config["tool_to_use_list"]) as mcp_tools:
        # Initialize the model
        model=OpenAILike(id="qwen-max", 
                         api_key=config["llm_api_key"],
                         base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")
        
        # Initialize the agent
        agent = Agent(model=model,
                      tools=[mcp_tools],
                      instructions=dedent("""\
                          你是一个阿里云云计算专家，请根据用户的问题，使用MCP服务查询阿里云的云产品信息，给出详细的解释。
                          请使用中文回答
                      """),
                      markdown=True,
                      show_tool_calls=True)

        # Run the agent
        await agent.aprint_response(config["user_question"], stream=True)

# Example usage
if __name__ == "__main__":
    # 1. User question goes here
    user_question = "我在上海有哪些ECS实例?"

    # 2. Prepare arguments for the agent
    url = "https://openapi-mcp.cn-hangzhou.aliyuncs.com/accounts/1411741061209533/custom/ecs_mcp/id/vwAaigJjvaOkaqHf/mcp"     # Full ECS List
    headers = {'Authorization': f'Bearer {os.getenv("ALI_OPENAPI_ACCESS_TOKEN")}'}
    llm_api_key = os.getenv("DASHSCOPE_API_KEY")

    # 3. Get selected tools list by user question
    seleted_tools = get_selected_tools_list(url, headers, llm_api_key, user_question)
    # print(seleted_tools)

    # # 4. Setup config
    config = {
        "server_url": url,
        "llm_api_key": llm_api_key,
        "headers": headers,
        "user_question": user_question,
        "tool_to_use_list": seleted_tools
    }

    # 5. Query LLM
    asyncio.run(run_agent(config))