import logging
from transformers import pipeline
import asyncio

class TextGenerationHandler:
    def __init__(self):
        # Initialize the text generation pipeline
        self.pipe = pipeline(
            "text2text-generation",
            model="google/flan-t5-small",
            max_length=2048,  # Increase max length for longer responses
            num_return_sequences=1
        )
        self.logger = logging.getLogger(__name__)
    
    async def generate_response(self, prompt: str) -> str:
        """
        Generate a complete response using the T5 model pipeline
        
        Args:
            prompt (str): Input text to generate from
            
        Returns:
            str: Generated text output
        """
        try:
            # Break down the generation into sections for better coherence
            sections = [
                "1. Executive Summary",
                "2. Project Scope and Objectives",
                "3. Architecture Overview",
                "4. Component Design",
                "5. Security and Compliance",
                "6. Deployment Strategy",
                "7. Team Requirements",
                "8. Cost Estimates",
                "9. Project Timeline"
            ]
            
            complete_response = []
            for section in sections:
                section_prompt = f"{prompt}\nGenerate content for: {section}"
                self.logger.info(f"Generating section: {section}")
                self.logger.debug(f"Section prompt: {section_prompt}")
                
                output = self.pipe(
                    section_prompt,
                    max_length=512,
                    do_sample=True,
                    temperature=0.7,
                    repetition_penalty=1.2,
                    no_repeat_ngram_size=3
                )
                
                section_text = output[0]['generated_text'].strip()
                self.logger.info(f"Generated text for {section}:\n{section_text}\n")
                
                complete_response.append(f"{section}\n{section_text}")
            
            final_response = "\n\n".join(complete_response)
            self.logger.info(f"Complete response:\n{final_response}")
            return final_response
            
        except Exception as e:
            self.logger.error(f"Error generating text: {str(e)}", exc_info=True)
            raise

    async def stream_response(self, prompt: str):
        """
        Stream the generated response section by section
        
        Args:
            prompt (str): Input text to generate from
            
        Yields:
            dict: Response chunks with type and content
        """
        try:
            # Generate complete response first
            response = await self.generate_response(prompt)
            
            # Stream each section
            accumulated_response = ""
            sections = response.split('\n\n')
            
            for section in sections:
                accumulated_response += section + "\n\n"
                self.logger.debug(f"Streaming section:\n{section}\n")
                
                yield {
                    "type": "content",
                    "content": accumulated_response.strip()
                }
                await asyncio.sleep(0.1)
                    
        except Exception as e:
            self.logger.error(f"Error in stream_response: {str(e)}", exc_info=True)
            yield {
                "type": "error",
                "content": str(e)
            }

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('f5_model.log')
    ]
)

f5_model = TextGenerationHandler()