import boto3
import json
import os
from typing import Optional, List, Dict, Any, Generator, TYPE_CHECKING
from pydantic import BaseModel
from .models.base import InfraConfig
from .features.infra.infra_config import InfraConfig as InfraConfigManager
import yaml
import time

if TYPE_CHECKING:
    from mypy_boto3_sqs import SQSClient, SQSResource
    from mypy_boto3_sqs.type_defs import MessageTypeDef
    from boto3.resources.base import ServiceResource

class SQSQueueConfig(InfraConfig):
    """
    Infrastructure configuration for an SQS queue.
    
    This class serves as the primary entry point for accessing SQS queues.
    It handles infrastructure configuration resolution and provides direct
    access to boto3 queue resources.
    """
    name: str
    fifo: bool = False
    
    def _resolve_queue_url(self, session: boto3.Session) -> str:
        """
        Resolves queue URL from infra.yml and validates it exists in AWS.
        """
        infra_config = InfraConfigManager.load()
        # Construct the output key from the resource name
        output_key = f"{self.name}QueueArn"
        queue_arn = infra_config.exports[output_key]
        
        # Extract queue name from ARN and construct URL
        arn_parts = queue_arn.split(':')
        extracted_queue_name = arn_parts[5]
        sqs_client = session.client('sqs')
        
        try:
            response = sqs_client.get_queue_url(QueueName=extracted_queue_name)
            return response['QueueUrl']
        except sqs_client.exceptions.QueueDoesNotExist:
            raise ValueError(f"Queue '{extracted_queue_name}' (from ARN in infra.yml) does not exist in AWS.")
    
    def get(self, session: Optional[boto3.Session] = None) -> 'SQSResource':
        """
        Returns the boto3 SQS queue resource for this queue configuration.
        
        Args:
            session: Optional boto3 session. If not provided, creates a new one.
            
        Returns:
            boto3 SQS queue resource object
        """
        session = session or boto3.Session()
        queue_url = self._resolve_queue_url(session)
        return session.resource('sqs').Queue(queue_url)
    
    def get_client(self, session: Optional[boto3.Session] = None) -> 'SQSClient':
        """
        Returns the boto3 SQS client.
        
        Args:
            session: Optional boto3 session. If not provided, creates a new one.
            
        Returns:
            boto3 SQS client object
        """
        session = session or boto3.Session()
        return session.client('sqs')


def process_queue_messages(queue: 'ServiceResource', batch_size: int = 1, poll_interval: int = 5) -> Generator['MessageTypeDef', None, None]:
    """
    Utility function to process messages from an SQS queue.
    
    This generator continuously polls the queue for messages and yields them one by one.
    Messages are automatically deleted after being yielded.
    
    Args:
        queue: boto3 SQS queue resource object
        batch_size: Maximum number of messages to retrieve per poll
        poll_interval: Time in seconds to wait between polls
        
    Yields:
        SQS message objects from the queue
    """
    while True:
        try:
            messages = queue.receive_messages(
                MaxNumberOfMessages=batch_size,
                WaitTimeSeconds=poll_interval,
                MessageAttributeNames=['All']
            )
            
            for message in messages:
                try:
                    yield message
                except Exception as e:
                    print(f"Error processing message: {e}")
                finally:
                    # Message is automatically deleted when it goes out of scope
                    # in boto3 resource API, but we can be explicit
                    message.delete()
                    
        except Exception as e:
            print(f"Error receiving messages: {e}")
            time.sleep(poll_interval)
