"""
{{ agent_name }} Agent
Generated by ABI-Core scaffolding
"""

import logging
import os
from typing import Dict, Any, Optional, AsyncIterable
from collections.abc import AsyncIterable

from a2a.types import (
    SendStreamingMessageSuccessResponse,
    TaskArtifactUpdateEvent,
    TaskState,
    TaskStatusUpdateEvent
)
from abi_core.common import prompts
from abi_core.common.utils import abi_logging
from abi_core.agent.agent import AbiAgent
from models import {{ agent_name.title().replace('_', '') }}Response

from langchain_core.messages import AIMessage
from langchain_ollama import ChatOllama
from langgraph.checkpoint.memory import MemorySaver
from langchain.agents import create_agent

logger = logging.getLogger(__name__)

MODEL_NAME = os.getenv('MODEL_NAME', '{{ model_name | default('qwen2.5:3b') }}')
memory = MemorySaver()

class {{ agent_class_name }}(AbiAgent):
    """{{ agent_description | default('ABI Agent for specialized tasks') }}"""
    
    def __init__(self):
        super().__init__(
            agent_name='{{ agent_name }}',
            description='{{ agent_description | default('Specialized ABI agent') }}',
            content_types=['text', 'text/plain', 'application/json']
        )
        
        # Initialize LLM
        self.llm = ChatOllama(
            model=MODEL_NAME,
            base_url=os.getenv('OLLAMA_HOST', 'http://localhost:11434'),
            temperature={{ temperature | default('0.1') }}
        )
        
        # Initialize agent
        self.agent = create_agent(
            self.llm,
            tools=[],  # Add tools as needed
            checkpointer=memory
        )
        
        abi_logging(f'[🚀] Starting ABI {{ agent_name }}')
    
    async def stream(
        self, 
        query: str, 
        context_id: str, 
        task_id: str
    ) -> AsyncIterable[Dict[str, Any]]:
        """
        Process query and stream responses
        
        Args:
            query: User query to process
            context_id: Context identifier
            task_id: Task identifier
            
        Yields:
            Dict containing response data
        """
        abi_logging(f'{{ agent_name }} processing query: {query[:100]}...')
        
        try:
            # Prepare agent configuration
            config = {
                "configurable": {
                    "thread_id": f"{context_id}-{task_id}"
                }
            }
            
            # Process with agent
            response_chunks = []
            async for chunk in self.agent.astream(
                {"messages": [("user", query)]}, 
                config=config
            ):
                if "agent" in chunk:
                    response_chunks.append(chunk["agent"]["messages"][0])
                elif "tools" in chunk:
                    # Handle tool responses
                    for tool_msg in chunk["tools"]["messages"]:
                        abi_logging(f'Tool response: {tool_msg.content[:100]}...')
            
            # Combine response
            if response_chunks:
                final_response = response_chunks[-1]
                content = final_response.content if hasattr(final_response, 'content') else str(final_response)
                
                yield {
                    'content': content,
                    'response_type': 'text',
                    'is_task_completed': True,
                    'require_user_input': False,
                    'metadata': {
                        'agent': '{{ agent_name }}',
                        'model': MODEL_NAME,
                        'context_id': context_id,
                        'task_id': task_id
                    }
                }
            else:
                yield {
                    'content': 'No response generated',
                    'response_type': 'text', 
                    'is_task_completed': True,
                    'require_user_input': False
                }
                
        except Exception as e:
            logger.error(f'Error in {{ agent_name }}: {e}')
            yield {
                'content': f'Error processing request: {str(e)}',
                'response_type': 'text',
                'is_task_completed': True,
                'require_user_input': False
            }

# Agent class is ready to be instantiated by the factory function