"""
Enhanced multi-agent system with dynamic marks configuration
"""

import json
import time
from typing import Dict, Any, Tuple
from agents import GeneratorAgent, VerifierAgent, FormatterAgent, SearchAgent
from prompts import get_generator_prompt, VERIFIER_PROMPT, FORMATTER_PROMPT

class MultiAgentSystem:
    """Enhanced orchestrator with dynamic marks support"""
    
    def __init__(self):
        self.generator = GeneratorAgent()
        self.verifier = VerifierAgent()
        self.formatter = FormatterAgent()
        self.searcher = SearchAgent()
        
        self.progress = {
            "stage": "ready",
            "message": "System ready",
            "percentage": 0
        }
    
    def update_progress(self, stage: str, message: str, percentage: int):
        """Update generation progress"""
        self.progress = {
            "stage": stage,
            "message": message,
            "percentage": percentage
        }
        print(f"Progress: {stage} - {message} ({percentage}%)")
    
    def extract_syllabus_from_pdf(self, pdf_file) -> str:
        """Extract text from uploaded syllabus PDF"""
        try:
            import PyPDF2
            pdf_reader = PyPDF2.PdfReader(pdf_file)
            text = ""
            for page in pdf_reader.pages:
                text += page.extract_text()
            return text[:3000]
        except Exception as e:
            print(f"PDF extraction error: {e}")
            return "Syllabus text extraction failed. Using default template."
    
    def extract_reference_from_pdf(self, pdf_file) -> str:
        """Extract text from reference question paper PDF"""
        if pdf_file is None:
            return "No reference question paper provided."
        
        try:
            import PyPDF2
            pdf_reader = PyPDF2.PdfReader(pdf_file)
            text = ""
            for page in pdf_reader.pages[:3]:
                text += page.extract_text()
            return text[:2000]
        except Exception as e:
            print(f"Reference PDF extraction error: {e}")
            return "Reference paper extraction failed."
    
    def generate_exam_package(self, 
                            subject: str,
                            stream: str,
                            part_a_count: int,
                            part_b_count: int, 
                            part_c_count: int,
                            part_a_marks: int = 2,
                            part_b_marks: int = 13,
                            part_c_marks: int = 14,
                            syllabus_file = None,
                            reference_file = None) -> Tuple[Dict[str, Any], str]:
        """Enhanced main method with dynamic marks"""
        
        try:
            # Stage 1: Data Preparation
            self.update_progress("preparation", "Extracting syllabus and reference data", 10)
            syllabus_text = self.extract_syllabus_from_pdf(syllabus_file)
            reference_text = self.extract_reference_from_pdf(reference_file)
            
            # Stage 2: Realtime Search
            self.update_progress("search", "Fetching recent developments", 20)
            realtime_updates = self.searcher.get_realtime_updates(subject)
            
            # Stage 3: Generation with dynamic marks
            self.update_progress("generation", "Generating question paper structure", 40)
            generator_prompt = get_generator_prompt(
                subject=subject,
                stream=stream,
                syllabus_text=syllabus_text,
                reference_text=reference_text,
                realtime_updates=realtime_updates,
                part_a_count=part_a_count,
                part_b_count=part_b_count,
                part_c_count=part_c_count,
                part_a_marks=part_a_marks,
                part_b_marks=part_b_marks,
                part_c_marks=part_c_marks
            )
            
            generated_content = self.generator.generate_question_paper(generator_prompt)
            
            if "error" in generated_content:
                return {}, "Generation failed: " + generated_content["error"]
            
            # Stage 4: Verification with dynamic marks
            self.update_progress("verification", "Verifying quality and standards", 60)
            verifier_prompt = VERIFIER_PROMPT.format(
                generated_content=json.dumps(generated_content, indent=2),
                bloom_mix="60% R/U, 40% A/An/Ev" if stream == "CSE" else "50% R/U, 50% A/An/Ev",
                part_a_count=part_a_count,
                part_b_count=part_b_count, 
                part_c_count=part_c_count,
                part_a_marks=part_a_marks,
                part_b_marks=part_b_marks,
                part_c_marks=part_c_marks,
                tag_requirements="Company tags" if stream == "CSE" else "GATE tags"
            )
            
            verification_result = self.verifier.verify_content(verifier_prompt)
            
            # Stage 5: Formatting
            self.update_progress("formatting", "Creating final outputs and answers", 80)
            formatter_prompt = FORMATTER_PROMPT.format(
                original_content=json.dumps(generated_content, indent=2),
                corrections=json.dumps(verification_result, indent=2)
            )
            
            final_output = self.formatter.format_final_output(formatter_prompt)
            
            # Stage 6: Completion
            self.update_progress("completion", "Package generation complete", 100)
            
            return final_output, "Success"
            
        except Exception as e:
            error_msg = f"System error: {str(e)}"
            print(error_msg)
            return {}, error_msg
    
    def get_progress(self) -> Dict[str, Any]:
        """Get current progress"""
        return self.progress