import json  # Import the json module

from qwen_agent.agents import Assistant

# Define LLM
llm_cfg = {
    'model': 'qwen3:0.6B',

    # Use the endpoint provided by Alibaba Model Studio:
    # 'model_type': 'qwen_dashscope',
    # 'api_key': os.getenv('DASHSCOPE_API_KEY'),

    # Use a custom endpoint compatible with OpenAI API:
    'model_server': 'http://localhost:11434/v1',  # api_base
    'api_key': 'EMPTY',

    # Other parameters:
    # 'generate_cfg': {
    #         # Add: When the response content is `<think>this is the thought</think>this is the answer;
    #         # Do not add: When the response has been separated by reasoning_content and content.
    #         'thought_in_content': True,
    #     },
}

# Define Tools
tools = [
    {'mcpServers': {  # You can specify the MCP configuration file
        'time': {
            'command': 'uvx',
            'args': ['mcp-server-time', '--local-timezone=Asia/Shanghai']
        },
        "fetch": {
            "command": "uvx",
            "args": ["mcp-server-fetch"]
        },
          "ddg-search": {
          "command": "npx",
          "args": ["-y", "@oevortex/ddg_search"]
        },
    }
    },
    'code_interpreter',  # Built-in tools
]

# Define Agent
bot = Assistant(llm=llm_cfg, function_list=tools)

# Streaming generation
messages = [{'role': 'user',
             'content':
                 """"
                  Write a 500 word blog post about the latest qwen 3 model. 
                  Use the search tool, and fetch the top 3 articles before you write the post. 
                  Write in a casual, but factual style - no hyperbole. 
                  Provide urls to the webpages in the output."""}]

final_responses = None
# Add comprehensive error handling around bot.run
try:
    timeout = 300  # 5 minute timeout
    final_responses = None
    partial_responses = []

    for responses in bot.run(messages=messages, enable_thinking=True, max_tokens=30000, timeout=timeout):
        print(".", end="", flush=True)
        partial_responses.append(responses)
        final_responses = responses.pop()

except TimeoutError:
    print("\nRequest timed out after {} seconds".format(timeout))
except ConnectionError as ce:
    print(f"\nConnection error occurred: {ce}")
except ValueError as ve:
    print(f"\nInvalid input or response format: {ve}")
except Exception as e:
    print(f"\nAn unexpected error occurred: {e}")
finally:
    if partial_responses and not final_responses:
        try:
            final_responses = partial_responses[-1].pop()
        except (IndexError, AttributeError):
            print("\nFailed to recover partial response")

# Pretty-print the final response object with validation
if final_responses:
    print("\n--- Full Response Object ---")
    try:
        if isinstance(final_responses, dict):
            print(json.dumps(final_responses, indent=2))
        else:
            print("Response is not in expected format:", type(final_responses))

        print("\n--- Extracted Content ---")
        if isinstance(final_responses, dict):
            content = final_responses.get('content')
            if content:
                print(content)
            else:
                print("No content field found in response")
        else:
            print("Cannot extract content from invalid response format")
    except json.JSONDecodeError as je:
        print(f"Failed to format response: {je}")
else:
    print("\nNo final response received from the agent.")