from io import BytesIO
from zipfile import ZipFile
from typing import Dict, Callable, Any, List
from hypotheses_generation import generate_hypotheses_from_system_charts

class HypothesesService:
    """
    Service class for Scientific Hypotheses Generation.
    Handles the orchestration of parsing system charts, generating hypotheses,
    and packaging results using the new synergistic approach.
    """
    
    def __init__(self):
        pass
    
    async def generate_hypotheses(
        self,
        system_charts_content_list: List[str],
        llm,
        material_goal_1: str = "Property #1 for xxx material",
        material_goal_2: str = "Property #2 for xxx material",
        iterations_count: int = 1,
        progress_callback: Callable[[Dict[str, Any]], None] = None
    ) -> Dict[str, Any]:
        """
        Main hypotheses generation pipeline using synergistic approach.
        
        Args:
            system_charts_content_list: List of content strings from multiple system chart files
            llm: Language model for hypothesis generation
            material_goal_1: First material design goal to guide hypothesis generation
            material_goal_2: Second material design goal to guide hypothesis generation
            iterations_count: Number of hypothesis iterations to generate
            progress_callback: Function to report progress updates
            
        Returns:
            Dict containing generation results and metadata
        """
        try:
            # Initial progress update
            if progress_callback:
                progress_callback({
                    "status": "starting",
                    "progress": 0,
                    "message": "Preparing for synergistic hypothesis generation..."
                })
            
            # Handle both single content and list of contents for backward compatibility
            if isinstance(system_charts_content_list, str):
                system_charts_content_list = [system_charts_content_list]
            
            # Create chart data structure for multiple files
            chart_data = []
            for i, content in enumerate(system_charts_content_list, 1):
                if content and content.strip():
                    chart_data.append({
                        'content': content,
                        'citation': f'System chart file {i}',
                        'title': f'System Chart File {i}',
                        'raw_content': content  # Pass the raw content
                    })
            
            if not chart_data:
                return {
                    "success": False,
                    "message": "No valid system chart content provided for hypothesis generation.",
                    "zip_data": None,
                    "filename": None
                }
            
            print(f"📊 Prepared {len(chart_data)} system chart files for synergistic hypothesis generation")
            
            # Update progress after parsing
            if progress_callback:
                progress_callback({
                    "status": "parsing_complete",
                    "progress": 10,
                    "total_charts": len(chart_data),
                    "message": f"System charts prepared from {len(chart_data)} files. Starting synergistic hypothesis generation..."
                })
            
            # Create progress callback function for the generation process
            def generation_progress_callback(progress_data):
                if progress_callback:
                    # Adjust progress to fit within our overall range (10% to 95%)
                    adjusted_progress = 10 + int((progress_data.get("progress", 0) / 100) * 85)
                    progress_data["progress"] = adjusted_progress
                    progress_callback(progress_data)
            
            # Generate hypotheses using the new synergistic approach
            print(f"🧠 Starting synergistic hypothesis generation with goals: '{material_goal_1}' and '{material_goal_2}' ({iterations_count} iteration{'s' if iterations_count > 1 else ''})...")
            hypotheses_content, hypotheses = await generate_hypotheses_from_system_charts(
                chart_data, 
                llm, 
                material_goal_1,
                material_goal_2,
                iterations_count,
                generation_progress_callback
            )
            
            # Note: Cross-chart analysis is now integrated into the main generation process
            cross_analysis_content = ""  # No longer separate
            
            # Create ZIP file in memory
            if progress_callback:
                progress_callback({
                    "status": "packaging",
                    "progress": 95,
                    "message": "Creating results package..."
                })
            
            zip_data = self._create_zip_file(
                hypotheses_content, 
                cross_analysis_content
            )
            
            # Prepare result
            return {
                "success": True,
                "message": f"Generated {iterations_count} iteration{'s' if iterations_count > 1 else ''} of synergistic hypotheses from {len(chart_data)} system chart files using cross-paper analysis approach!",
                "zip_data": zip_data,
                "filename": "hypotheses.zip",
                "hypotheses": hypotheses,
                "charts_processed": len(chart_data),
                "cross_analysis_included": True,  # Now always included in main process
                "method": "synergistic_analysis",
                "iterations_count": iterations_count
            }
            
        except Exception as e:
            return {
                "success": False,
                "message": f"Error during synergistic hypothesis generation: {str(e)}",
                "zip_data": None,
                "filename": None
            }
    
    def _create_zip_file(
        self, 
        hypotheses_content: str, 
        cross_analysis_content: str = ""
    ) -> bytes:
        """Create ZIP file containing hypothesis generation results."""
        zip_buffer = BytesIO()
        
        with ZipFile(zip_buffer, "w") as zipf:
            # Add main hypotheses file
            zipf.writestr("hypotheses.txt", hypotheses_content)

            print(f"📦 Created synergistic hypotheses ZIP package with {len(zipf.namelist())} files")
        
        # Get ZIP data and close buffer
        zip_data = zip_buffer.getvalue()
        zip_buffer.close()
        
        return zip_data
    
    @staticmethod
    def validate_multiple_system_charts_content(content_list: List[str]) -> Dict[str, Any]:
        """
        Validate that the provided content list is suitable for synergistic hypothesis generation.
        
        Args:
            content_list: List of content strings to validate
            
        Returns:
            Dict with validation results
        """
        try:
            if not content_list:
                return {
                    "valid": False,
                    "message": "No content provided for hypothesis generation.",
                    "chart_count": 0
                }
            
            valid_files = 0
            total_data_rows = 0
            file_details = []
            
            for i, content in enumerate(content_list, 1):
                if not content or not content.strip():
                    file_details.append(f"File {i}: Empty content")
                    continue
                
                # Check for new format with 100 '=' delimiters
                lines = content.strip().split('\n')
                has_new_format = any(line.strip() == '=' * 100 for line in lines)
                has_paper_keywords = any(line.strip().startswith('PAPER:') for line in lines)
                has_chart_keywords = any(line.strip().startswith('SYSTEM_CHART:') for line in lines)
                
                if has_new_format and has_paper_keywords and has_chart_keywords:
                    # Count entries in new format
                    paper_count = len([line for line in lines if line.strip().startswith('PAPER:')])
                    chart_count = len([line for line in lines if line.strip().startswith('SYSTEM_CHART:')])
                    
                    if paper_count > 0 and chart_count > 0 and paper_count == chart_count:
                        valid_files += 1
                        total_data_rows += paper_count
                        file_details.append(f"File {i}: Valid format with {paper_count} entries")
                    else:
                        file_details.append(f"File {i}: Format issues - {paper_count} papers, {chart_count} charts")
                else:
                    # Fallback to table-based validation
                    table_lines = [line for line in lines if '|' in line and not line.strip().startswith('|---')]
                    
                    if table_lines:
                        data_rows = len([line for line in table_lines if line.strip() and (
                            ('|' in line) and not any(
                                header_word in line.lower() for header_word in ['processing', 'structure', 'property', 'mechanism']
                            )
                        )])
                        
                        if data_rows >= 2:
                            valid_files += 1
                            total_data_rows += data_rows
                            file_details.append(f"File {i}: Valid table format with {data_rows} rows")
                        else:
                            file_details.append(f"File {i}: Insufficient table data ({data_rows} rows)")
                    else:
                        file_details.append(f"File {i}: No valid format found")
            
            if valid_files == 0:
                return {
                    "valid": False,
                    "message": f"No valid system chart files found. Details: {'; '.join(file_details)}",
                    "chart_count": 0,
                    "valid_charts": 0,
                    "file_details": file_details
                }
            
            if total_data_rows < 2:
                return {
                    "valid": False,
                    "message": f"Insufficient data across all files. Need multiple entries for synergistic analysis. Found {total_data_rows} total entries.",
                    "chart_count": len(content_list),
                    "valid_charts": valid_files,
                    "file_details": file_details
                }
            
            return {
                "valid": True,
                "message": f"Validated {valid_files}/{len(content_list)} system chart files for synergistic hypothesis generation. Found {total_data_rows} total entries.",
                "chart_count": len(content_list),
                "valid_charts": valid_files,
                "data_entries": total_data_rows,
                "file_details": file_details
            }
            
        except Exception as e:
            return {
                "valid": False,
                "message": f"Error validating content: {str(e)}",
                "chart_count": 0,
                "file_details": []
            }

    @staticmethod
    def validate_system_charts_content(content) -> Dict[str, Any]:
        """
        Validate that the provided content is suitable for synergistic hypothesis generation.
        Handles both single content string and list of content strings.
        
        Args:
            content: Content to validate (str or List[str])
            
        Returns:
            Dict with validation results
        """
        # Handle both single content and list of contents
        if isinstance(content, str):
            return HypothesesService.validate_multiple_system_charts_content([content])
        elif isinstance(content, list):
            return HypothesesService.validate_multiple_system_charts_content(content)
        else:
            return {
                "valid": False,
                "message": "Invalid content type. Expected string or list of strings.",
                "chart_count": 0
            }