# agents/report_generation_agent.py
import logging
from typing import Dict, List, Optional, Tuple, Union, Any
import json
import time
from datetime import datetime

# Import latest LangChain packages
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
#from langchain_core.pydantic_v1 import BaseModel, Field
from pydantic import BaseModel, Field 

from langchain_community.llms.huggingface_text_gen_inference import HuggingFaceTextGenInference
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
#from langchain_community.llms import HuggingFaceHub
#from langchain_huggingface import HuggingFaceHub
from langchain_community.llms import HuggingFacePipeline
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

class ReportStructure(BaseModel):
    """Structure for the generated report."""
    executive_summary: str = Field(description="Concise summary of key findings with confidence level")
    topic_overview: str = Field(description="Brief introduction to the topic")
    text_analysis: str = Field(description="Summary of relevant text findings")
    image_analysis: str = Field(description="Summary of relevant image findings")
    confidence_assessment: str = Field(description="Explanation of confidence level and evidence quality")
    detailed_findings: str = Field(description="Comprehensive analysis of all relevant information")
    conclusion: str = Field(description="Final insights and potential next steps")

class ReportGeneratorAgent:
    def __init__(self, summary_model_manager=None, token_manager=None, 
                 cache_manager=None, metrics_calculator=None):
        """Initialize the ReportGeneratorAgent with required utilities."""
        self.logger = logging.getLogger(__name__)
        self.summary_model_manager = summary_model_manager
        self.token_manager = token_manager
        self.cache_manager = cache_manager
        self.metrics_calculator = metrics_calculator
        
        # Agent name for logging
        self.agent_name = "report_generation_agent"
        
        # Initialize LangChain components
        self._initialize_langchain_components()
        
    def _initialize_langchain_components(self):
        """Initialize LangChain components for report generation."""
        try:
            # Use HuggingFaceHub with a local model that doesn't require API keys
            # We'll use a smaller model since we're running locally
            # self.llm = HuggingFaceHub(
            #     repo_id="google/flan-t5-small",
            #     model_kwargs={"temperature": 0.7, "max_length": 1024}
            # )
            tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
            model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
            pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer, max_length=1024)
            self.llm = HuggingFacePipeline(pipeline=pipe)
            
            # Create prompt templates
            self._create_prompt_templates()
            
            # Create LangChain chains
            self._create_langchain_chains()
            
            self.logger.info("LangChain components initialized successfully")
        except Exception as e:
            self.logger.error(f"Failed to initialize LangChain components: {e}")
            # Fallback to summary model manager if LangChain initialization fails
            self.llm = None
    
    def _create_prompt_templates(self):
        """Create prompt templates for different report generation tasks."""
        # Executive summary prompt
        self.executive_summary_prompt = PromptTemplate.from_template(
            """
            Generate an executive summary for a report on the topic: {topic}
            
            The overall confidence level is: {confidence_level}
            
            Text analysis found {text_count} relevant documents.
            Image analysis found {image_count} relevant images.
            
            Key text findings:
            {text_findings}
            
            Key image findings:
            {image_findings}
            
            Create a concise, professional executive summary that clearly communicates:
            1. The main findings related to the topic
            2. The confidence level in these findings
            3. The strength of the evidence
            
            Executive Summary:
            """
        )
        
        # Detailed report prompt
        self.detailed_report_prompt = ChatPromptTemplate.from_messages([
            ("system", """You are an expert report generator that synthesizes information from multiple sources.
            Your task is to create a comprehensive, well-structured report based on text and image analyses.
            Adjust your level of detail and certainty based on the confidence level.
            For high confidence, be definitive. For medium confidence, be more measured. For low confidence, be appropriately cautious.
            """),
            ("user", """
            Topic: {topic}
            
            Overall Confidence Level: {confidence_level}
            
            Text Analysis:
            {text_analysis}
            
            Image Analysis:
            {image_analysis}
            
            Please generate a complete report with the following sections:
            1. Executive Summary
            2. Topic Overview
            3. Text Analysis Findings
            4. Image Analysis Findings
            5. Confidence Assessment
            6. Detailed Findings
            7. Conclusion
            
            Format the report in markdown with appropriate headings and structure.
            """)
        ])
    
    def _create_langchain_chains(self):
        """Create LangChain chains for report generation."""
        # Executive summary chain
        self.executive_summary_chain = (
            self.executive_summary_prompt 
            | self.llm 
            | StrOutputParser()
        )
        
        # Detailed report chain
        self.detailed_report_chain = (
            self.detailed_report_prompt 
            | self.llm 
            | StrOutputParser()
        )
    
    def _prepare_input_data(self, topic: str, text_analysis: Dict[str, Any], 
                           image_analysis: Dict[str, Any]) -> Dict[str, Any]:
        """Prepare input data for report generation."""
        # Extract text findings
        text_findings = ""
        if text_analysis and "document_analyses" in text_analysis:
            for i, doc in enumerate(text_analysis.get("document_analyses", [])[:3]):  # Top 3 docs
                text_findings += f"- Document {i+1}: {doc.get('summary', 'No summary')}.\n"
        
        # Extract image findings
        image_findings = ""
        if image_analysis and "image_analyses" in image_analysis:
            for i, img in enumerate(image_analysis.get("image_analyses", [])[:3]):  # Top 3 images
                image_findings += f"- Image {i+1}: {img.get('caption', 'No caption')}.\n"
        
        # Determine overall confidence
        text_confidence = text_analysis.get("overall_confidence", 0) if text_analysis else 0
        image_confidence = image_analysis.get("overall_confidence", 0) if image_analysis else 0
        
        # Weight text more heavily (70/30 split)
        if text_analysis and image_analysis:
            overall_confidence = 0.7 * text_confidence + 0.3 * image_confidence
        elif text_analysis:
            overall_confidence = text_confidence
        elif image_analysis:
            overall_confidence = image_confidence
        else:
            overall_confidence = 0
            
        # Map numerical confidence to level
        if overall_confidence >= 0.7:
            confidence_level = "high"
        elif overall_confidence >= 0.4:
            confidence_level = "medium"
        else:
            confidence_level = "low"
            
        # Prepare complete text analysis for detailed report
        full_text_analysis = "No text analysis available."
        if text_analysis:
            full_text_analysis = f"""
            {text_analysis.get('relevant_documents', 0)} relevant documents found out of {text_analysis.get('total_documents', 0)}.
            Confidence level: {text_analysis.get('confidence_level', 'unknown')}.
            
            Document findings:
            """
            for i, doc in enumerate(text_analysis.get("document_analyses", [])):
                full_text_analysis += f"\n{i+1}. {doc.get('filename', 'Unknown document')}: {doc.get('summary', 'No summary')}"
        
        # Prepare complete image analysis for detailed report
        full_image_analysis = "No image analysis available."
        if image_analysis:
            full_image_analysis = f"""
            {image_analysis.get('relevant_images', 0)} relevant images found out of {image_analysis.get('total_images', 0)}.
            Confidence level: {image_analysis.get('confidence_level', 'unknown')}.
            
            Image findings:
            """
            for i, img in enumerate(image_analysis.get("image_analyses", [])):
                full_image_analysis += f"\n{i+1}. {img.get('filename', 'Unknown image')}: {img.get('caption', 'No caption')}"
        
        return {
            "topic": topic,
            "confidence_level": confidence_level,
            "text_count": text_analysis.get("relevant_documents", 0) if text_analysis else 0,
            "image_count": image_analysis.get("relevant_images", 0) if image_analysis else 0,
            "text_findings": text_findings,
            "image_findings": image_findings,
            "text_analysis": full_text_analysis,
            "image_analysis": full_image_analysis,
            "overall_confidence": overall_confidence
        }
    
    def generate_report(self, topic: str, text_analysis: Dict[str, Any], 
                       image_analysis: Dict[str, Any]) -> Dict[str, Any]:
        """
        Generate a comprehensive report based on text and image analyses.
        Returns the report and metadata.
        """
        start_time = time.time()
        self.logger.info(f"Generating report for topic: {topic}")
        
        # Check if we can use the summary model directly
        if self.summary_model_manager and not self.llm:
            return self._generate_report_with_summary_model(topic, text_analysis, image_analysis)
        
        # Prepare input data
        input_data = self._prepare_input_data(topic, text_analysis, image_analysis)
        
        try:
            # Generate executive summary
            executive_summary = self.executive_summary_chain.invoke(input_data)
            
            # Track token usage if available
            if self.token_manager:
                # Estimate token count (approximate)
                summary_tokens = len(executive_summary.split()) * 1.3  # Rough estimate
                self.token_manager.log_usage(
                    self.agent_name, "report_generation", int(summary_tokens), "langchain")
                
                # Log energy usage if metrics calculator is available
                if self.metrics_calculator:
                    energy_usage = self.token_manager.calculate_energy_usage(
                        int(summary_tokens), "langchain")
                    self.metrics_calculator.log_energy_usage(
                        energy_usage, "langchain", self.agent_name, "report_generation")
            
            # Generate detailed report
            detailed_report = self.detailed_report_chain.invoke(input_data)
            
            # Track token usage for detailed report
            if self.token_manager:
                # Estimate token count (approximate)
                detailed_tokens = len(detailed_report.split()) * 1.3  # Rough estimate
                self.token_manager.log_usage(
                    self.agent_name, "report_generation", int(detailed_tokens), "langchain")
                
                # Log energy usage if metrics calculator is available
                if self.metrics_calculator:
                    energy_usage = self.token_manager.calculate_energy_usage(
                        int(detailed_tokens), "langchain")
                    self.metrics_calculator.log_energy_usage(
                        energy_usage, "langchain", self.agent_name, "report_generation")
            
            # Prepare final report
            report = {
                "topic": topic,
                "timestamp": datetime.now().isoformat(),
                "executive_summary": executive_summary,
                "detailed_report": detailed_report,
                "confidence_level": input_data["confidence_level"],
                "confidence_score": input_data["overall_confidence"],
                "sources": {
                    "text_documents": text_analysis.get("relevant_documents", 0) if text_analysis else 0,
                    "images": image_analysis.get("relevant_images", 0) if image_analysis else 0
                }
            }
            
            # Add processing metadata
            processing_time = time.time() - start_time
            report["processing_time"] = processing_time
            
            self.logger.info(f"Report generation completed in {processing_time:.2f} seconds.")
            
            return report
            
        except Exception as e:
            self.logger.error(f"Failed to generate report with LangChain: {e}")
            # Fallback to summary model
            return self._generate_report_with_summary_model(topic, text_analysis, image_analysis)
    
    def _generate_report_with_summary_model(self, topic: str, text_analysis: Dict[str, Any], 
                                          image_analysis: Dict[str, Any]) -> Dict[str, Any]:
        """Fallback method to generate report using the summary model manager."""
        self.logger.info("Using summary model fallback for report generation")
        
        if not self.summary_model_manager:
            return {
                "topic": topic,
                "error": "No report generation capability available",
                "timestamp": datetime.now().isoformat()
            }
        
        # Extract document analyses
        doc_analyses = []
        if text_analysis and "document_analyses" in text_analysis:
            doc_analyses = text_analysis.get("document_analyses", [])
        
        # Extract image analyses
        img_analyses = []
        if image_analysis and "image_analyses" in image_analysis:
            img_analyses = image_analysis.get("image_analyses", [])
        
        # Use summary model to combine analyses
        report = self.summary_model_manager.combine_analyses(
            doc_analyses, img_analyses, topic, self.agent_name)
        
        # Add timestamp
        report["timestamp"] = datetime.now().isoformat()
        
        return report
    
    def generate_confidence_statement(self, confidence_level: str) -> str:
        """Generate an appropriate confidence statement based on the level."""
        if confidence_level == "high":
            return "This analysis is provided with high confidence based on strong evidence in the provided materials."
        elif confidence_level == "medium":
            return "This analysis is provided with moderate confidence. Some aspects may require additional verification."
        else:
            return "This analysis is provided with low confidence due to limited relevant information in the provided materials."
    
    def get_cached_report(self, topic: str, text_analysis_id: str, image_analysis_id: str) -> Optional[Dict[str, Any]]:
        """
        Try to retrieve a previously generated report from cache.
        Returns None if not found in cache.
        """
        if not self.cache_manager:
            return None
        
        # Create a cache key based on inputs
        cache_key = f"report:{topic}:{text_analysis_id}:{image_analysis_id}"
        
        # Try to get from cache
        cache_hit, cached_report = self.cache_manager.get(cache_key, namespace="reports")
        
        if cache_hit and cached_report:
            # Update metrics if available
            if self.metrics_calculator:
                self.metrics_calculator.update_cache_metrics(1, 0, 0.02)  # Estimated energy saving
                self.metrics_calculator.log_tokens_saved(500)  # Approximate tokens saved
            
            self.logger.info(f"Retrieved cached report for topic: {topic}")
            return cached_report
        
        return None
    
    def format_report_for_display(self, report: Dict[str, Any], format: str = "markdown") -> str:
        """Format the report for display in the specified format."""
        if format == "markdown":
            # Format as markdown
            md = f"# Report: {report.get('topic', 'Unknown Topic')}\n\n"
            md += f"*Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*\n\n"
            
            md += f"## Executive Summary\n\n{report.get('executive_summary', 'No summary available.')}\n\n"
            
            md += f"**Confidence Level: {report.get('confidence_level', 'unknown').title()}**\n\n"
            md += f"*{self.generate_confidence_statement(report.get('confidence_level', 'low'))}*\n\n"
            
            md += f"## Detailed Report\n\n{report.get('detailed_report', 'No detailed report available.')}\n\n"
            
            md += f"## Sources\n\n"
            md += f"- Text Documents: {report.get('sources', {}).get('text_documents', 0)}\n"
            md += f"- Images: {report.get('sources', {}).get('images', 0)}\n"
            
            return md
            
        elif format == "html":
            # Format as HTML
            # html = fr"""
            # <div class="report">
            #     <h1>Report: {report.get('topic', 'Unknown Topic')}</h1>
            #     <p class="timestamp"><em>Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</em></p>
                
            #     <h2>Executive Summary</h2>
            #     <div class="summary">
            #         <p>{report.get('executive_summary', 'No summary available.')}</p>
            #     </div>
                
            #     <div class="confidence">
            #         <p><strong>Confidence Level: {report.get('confidence_level', 'unknown').title()}</strong></p>
            #         <p><em>{self.generate_confidence_statement(report.get('confidence_level', 'low'))}</em></p>
            #     </div>
                
            #     <h2>Detailed Report</h2>
            #     <div class="detailed-report">
            #         {report.get('detailed_report', 'No detailed report available.').replace('\\n', '<br>')}
            #     </div>
                
            #     <h2>Sources</h2>
            #     <ul>
            #         <li>Text Documents: {report.get('sources', {}).get('text_documents', 0)}</li>
            #         <li>Images: {report.get('sources', {}).get('images', 0)}</li>
            #     </ul>
            # </div>
            # """
            # Start with the opening tags
            html = "<div class=\"report\">\n"
            html += f"    <h1>Report: {report.get('topic', 'Unknown Topic')}</h1>\n"
            html += f"    <p class=\"timestamp\"><em>Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</em></p>\n"
                
            # Executive Summary section
            html += "    <h2>Executive Summary</h2>\n"
            html += "    <div class=\"summary\">\n"
            html += f"        <p>{report.get('executive_summary', 'No summary available.')}</p>\n"
            html += "    </div>\n"
                
            # Confidence section
            html += "    <div class=\"confidence\">\n"
            html += f"        <p><strong>Confidence Level: {report.get('confidence_level', 'unknown').title()}</strong></p>\n"
            html += f"        <p><em>{self.generate_confidence_statement(report.get('confidence_level', 'low'))}</em></p>\n"
            html += "    </div>\n"
                
            # Detailed Report section
            html += "    <h2>Detailed Report</h2>\n"
            html += "    <div class=\"detailed-report\">\n"
            detailed_report = report.get('detailed_report', 'No detailed report available.').replace('\n', '<br>')
            html += f"        {detailed_report}\n"
            html += "    </div>\n"
                
            # Sources section
            html += "    <h2>Sources</h2>\n"
            html += "    <ul>\n"
            html += f"        <li>Text Documents: {report.get('sources', {}).get('text_documents', 0)}</li>\n"
            html += f"        <li>Images: {report.get('sources', {}).get('images', 0)}</li>\n"
            html += "    </ul>\n"
            
            # Close the main div
            html += "</div>\n"

            return html
        
        else:
            return f"Unsupported format: {format}"