from agency_swarm.agents import Agent
from agency_swarm.tools import CodeInterpreter
import os
import logging
from .tools.SearchAndScrape import SearchAndScrape

class ValidationAgent(Agent):
    def __init__(self):
        super().__init__(
            name="ValidationAgent",
            description="This agent validates market research reports using AI and ensures data completeness.",
            instructions="./instructions.md",
            files_folder="./files",
            schemas_folder="./schemas",
            tools=[SearchAndScrape],
            tools_folder="./tools",
            temperature=0.3,
            model="groq/llama-3.3-70b-versatile", 
            max_prompt_tokens=25000,
        )

    def validate_data(self, report_data):
        """Validate the report using AI and fill gaps if needed."""
        validation_prompt = f"""
        Analyze this market research report for quality and completeness:

        {report_data}

        Please check for:
        1. Missing key information
        2. Data accuracy and consistency
        3. Logical flow and structure
        4. Completeness of sections:
           - Market Size & Growth
           - Competitive Landscape
           - Consumer Analysis
           - Technology & Innovation
           - Future Outlook

        Provide a detailed assessment with:
        1. Quality score (0-100)
        2. List of missing or incomplete sections
        3. Specific recommendations for improvement
        4. Additional data points needed

        Format: JSON with these keys:
        {
            "quality_score": int,
            "missing_sections": list,
            "recommendations": list,
            "additional_data_needed": list,
            "is_complete": boolean
        }
        """

        try:
            model = self._get_model()  # Updated method to get model
            response = model.generate_content(validation_prompt)
            validation_result = response.text
            
            # If validation shows missing data, scrape for it
            if '"is_complete": false' in validation_result.lower():
                missing_data = self._fill_missing_data(validation_result)
                if missing_data:
                    # Combine original report with new data
                    updated_report = self._merge_reports(report_data, missing_data)
                    # Validate again
                    return self.validate_data(updated_report)
            
            return validation_result

        except Exception as e:
            logging.error(f"Validation error: {str(e)}")
            return {"error": str(e)}

    def _fill_missing_data(self, validation_result):
        """Fill missing data based on validation results."""
        try:
            # Extract missing sections from validation result
            import json
            result = json.loads(validation_result)
            missing_sections = result.get("missing_sections", [])
            
            additional_data = []
            for section in missing_sections:
                # Create specific search query for missing section
                search_query = f"{section} market research data analysis"
                tool = SearchAndScrape(query=search_query)
                section_data = tool.run()
                if section_data:
                    additional_data.append({
                        "section": section,
                        "content": section_data
                    })
            
            return additional_data if additional_data else None

        except Exception as e:
            logging.error(f"Error filling missing data: {str(e)}")
            return None

    def _merge_reports(self, original_report, new_data):
        """Merge original report with newly scraped data."""
        merge_prompt = f"""
        Merge this original report with new data:

        Original Report:
        {original_report}

        New Data to Add:
        {new_data}

        Please create a cohesive, well-structured report that incorporates all information without duplication.
        Ensure proper flow and transitions between sections.
        """

        try:
            model = self._get_model()  # Updated method to get model
            response = model.generate_content(merge_prompt)
            return response.text
        except Exception as e:
            logging.error(f"Error merging reports: {str(e)}")
            return original_report

    def response_validator(self, message):
        return message