import json
import sys
import argparse
from typing import Optional, List, Dict, Any
from dataclasses import dataclass
from rich.console import Console
from rich.spinner import Spinner
from qwen_agent.agents import Assistant
from transformers import pipeline

@dataclass
class AgentConfig:
    model: str
    server: str
    api_key: str
    max_tokens: int = 30000
    enable_thinking: bool = True
    tools: Optional[List[Dict[str, Any]]] = None

    def __post_init__(self):
        if self.tools is None:
            self.tools = [
                {'mcpServers': {
                    'time': {
                        'command': 'uvx',
                        'args': ['mcp-server-time', '--local-timezone=Europe/London']
                    },
                    "fetch": {
                        "command": "uvx",
                        "args": ["mcp-server-fetch"]
                    },
                    "ddg-search": {
                        "command": "npx",
                        "args": ["-y", "duckduckgo-mcp-server"]
                    },
                }},
                'code_interpreter',
            ]

def setup_argparse():
    parser = argparse.ArgumentParser(description='Qwen3 Agent CLI')
    parser.add_argument('--model', default='qwen3:32b',
                       help='Model identifier (default: qwen3:32b)')
    parser.add_argument('--server', default='http://localhost:11434/v1',
                       help='Model server URL (default: http://localhost:11434/v1)')
    parser.add_argument('--api-key', default='EMPTY',
                       help='API key for the model server (default: EMPTY)')
    
    subparsers = parser.add_subparsers(dest='command', help='Available commands')
    
    # Prompt command
    prompt_parser = subparsers.add_parser('prompt', help='Run agent with a prompt')
    prompt_parser.add_argument('text', nargs='?', default='-',
                             help='Prompt text or "-" for stdin (default: -)')
    
    return parser

def read_prompt(text: str) -> str:
    """Read prompt from argument or stdin if text is '-'"""
    if text == '-':
        return sys.stdin.read().strip()
    return text

def run_agent(config: AgentConfig, prompt: str) -> None:
    """Run the agent with the given configuration and prompt"""
    llm_cfg = {
        'model': config.model,
        'model_server': config.server,
        'api_key': config.api_key,
    }

    # Define Agent
    bot = Assistant(llm=llm_cfg, function_list=config.tools)
    console = Console()

    # Streaming generation
    messages = [{'role': 'user', 'content': prompt}]

    final_responses = None
    try:
        with console.status("[bold blue]Thinking...", spinner="dots") as status:
            for responses in bot.run(messages=messages, 
                                   enable_thinking=config.enable_thinking, 
                                   max_tokens=config.max_tokens):
                final_responses = responses.pop()
    except Exception as e:
        console.print(f"[bold red]An error occurred during agent execution:[/] {e}")

    # Pretty-print the final response object
    if final_responses:
        console.print("\n[bold green]--- Full Response Object ---[/]")
        console.print(json.dumps(final_responses, indent=2))
        console.print("\n[bold green]--- Extracted Content ---[/]")
        console.print(final_responses.get('content', 'No content found in response.'))
    else:
        console.print("[bold red]No final response received from the agent.[/]")

def main():
    parser = setup_argparse()
    args = parser.parse_args()
    
    if args.command == 'prompt':
        prompt_text = read_prompt(args.text)
        config = AgentConfig(
            model=args.model,
            server=args.server,
            api_key=args.api_key
        )
        run_agent(config, prompt_text)
    else:
        parser.print_help()

if __name__ == '__main__':
    main()