from typing import Dict, List, Optional, Callable
from datetime import datetime
from .utilties.search_utilities import remove_think_tags, format_findings_to_text, format_links
import os
from .utilties.enums import KnowledgeAccumulationApproach
from .config import settings, get_llm, get_search 
from .citation_handler import CitationHandler
from datetime import datetime
from .utilties.search_utilities import extract_links_from_search_results
import logging
logger = logging.getLogger(__name__)
class AdvancedSearchSystem:
    def __init__(self):

        
        # Get fresh configuration

        self.search = get_search()
        self.model = get_llm()
        self.max_iterations = settings.search.iterations
        self.questions_per_iteration = settings.search.questions_per_iteration
        
        self.context_limit = settings.general.knowledge_accumulation_context_limit
        self.questions_by_iteration = {}
        self.citation_handler = CitationHandler(self.model)
        self.progress_callback = None
        self.all_links_of_system = list()
        
        # Check if search is available, log warning if not
        if self.search is None:
            logger.info("WARNING: Search system initialized with no search engine! Research will not be effective.")
            self._update_progress("WARNING: No search engine available", None, {"error": "No search engine configured properly"})

        

    def set_progress_callback(self, callback: Callable[[str, int, dict], None]) -> None:
        """Set a callback function to receive progress updates.
        
        Args:
            callback: Function that takes (message, progress_percent, metadata)
        """
        self.progress_callback = callback

    def _update_progress(self, message: str, progress_percent: int = None, metadata: dict = None) -> None:
        """Send a progress update via the callback if available.
        
        Args:
            message: Description of the current progress state
            progress_percent: Progress percentage (0-100), if applicable
            metadata: Additional data about the progress state
        """
        if self.progress_callback:
            self.progress_callback(message, progress_percent, metadata or {})

    def _get_follow_up_questions(self, current_knowledge: str, query: str) -> List[str]:
        now = datetime.now()
        current_time = now.strftime("%Y-%m-%d")
        
        self._update_progress("Generating follow-up questions...", None, {"iteration": len(self.questions_by_iteration)})
        
        if self.questions_by_iteration:
            prompt = f"""Critically reflect current knowledge (e.g., timeliness), what {self.questions_per_iteration} high-quality internet search questions remain unanswered to exactly answer the query?
            Query: {query}
            Today: {current_time} 
            Past questions: {str(self.questions_by_iteration)}
            Knowledge: {current_knowledge}
            Include questions that critically reflect current knowledge.
            \n\n\nFormat: One question per line, e.g. \n Q: question1 \n Q: question2\n\n"""
        else:
            prompt = f" You will have follow up questions. First, identify if your knowledge is outdated (high chance). Today: {current_time}. Generate {self.questions_per_iteration} high-quality internet search questions to exactly answer: {query}\n\n\nFormat: One question per line, e.g. \n Q: question1 \n Q: question2\n\n"

        response = self.model.invoke(prompt)
        questions = [
            q.replace("Q:", "").strip()
            for q in remove_think_tags(response.content).split("\n")
            if q.strip().startswith("Q:")
        ][: self.questions_per_iteration]
        
        self._update_progress(
            f"Generated {len(questions)} follow-up questions", 
            None, 
            {"questions": questions}
        )
        
        return questions

    def _compress_knowledge(self, current_knowledge: str, query: str, section_links) -> List[str]:
        self._update_progress("Compressing and summarizing knowledge...", None)

        now = datetime.now()
        current_time = now.strftime("%Y-%m-%d")
        formatted_links = format_links(links=section_links)
        if self.questions_by_iteration:
            prompt = f"""First provide a high-quality 1 page explanation with IEEE Referencing Style e.g. [1,2]. Never make up sources. Than provide a exact high-quality one sentence-long answer to the query. 

            Knowledge: {current_knowledge}
            Query: {query}
            I will append following text to your output for the sources (dont repeat it):\n\n {formatted_links}"""
        response = self.model.invoke(prompt)
        
        self._update_progress("Knowledge compression complete", None)
        response = remove_think_tags(response.content)
        response = str(response) #+ "\n\n" + str(formatted_links)

        return response

    def analyze_topic(self, query: str) -> Dict:
        logger.info(f"Starting research on topic: {query}")
        


        findings = []
        current_knowledge = ""
        iteration = 0
        total_iterations = self.max_iterations
        section_links = list()
        
        self._update_progress("Initializing research system", 5, {
            "phase": "init",
            "iterations_planned": total_iterations
        })
        
        # Check if search engine is available
        if self.search is None:
            error_msg = "Error: No search engine available. Please check your configuration."
            self._update_progress(error_msg, 100, {
                "phase": "error", 
                "error": "No search engine available",
                "status": "failed"
            })
            return {
                "findings": [],
                "iterations": 0,
                "questions": {},
                "formatted_findings": "Error: Unable to conduct research without a search engine.",
                "current_knowledge": "",
                "error": error_msg
            }

        while iteration < self.max_iterations:
            iteration_progress_base = (iteration / total_iterations) * 100
            self._update_progress(f"Starting iteration {iteration + 1} of {total_iterations}", 
                                 int(iteration_progress_base),
                                 {"phase": "iteration_start", "iteration": iteration + 1})
            
            # Generate questions for this iteration
            questions = self._get_follow_up_questions(current_knowledge, query)
            self.questions_by_iteration[iteration] = questions
            logger.info(f"Generated questions: {questions}")
            question_count = len(questions)
            for q_idx, question in enumerate(questions):
                question_progress_base = iteration_progress_base + (((q_idx+1) / question_count) * (100/total_iterations) * 0.5)
                
                self._update_progress(f"Searching for: {question}", 
                                     int(question_progress_base),
                                     {"phase": "search", "iteration": iteration + 1, "question_index": q_idx + 1})
                
                try:
                    if self.search is None:
                        self._update_progress(f"Search engine unavailable, skipping search for: {question}", 
                                            int(question_progress_base + 2),
                                            {"phase": "search_error", "error": "No search engine available"})
                        search_results = []
                    else:
                        search_results = self.search.run(question)
                except Exception as e:
                    error_msg = f"Error during search: {str(e)}"
                    logger.info(f"SEARCH ERROR: {error_msg}")
                    self._update_progress(error_msg, 
                                        int(question_progress_base + 2),
                                        {"phase": "search_error", "error": str(e)})
                    search_results = []
                
                if search_results is None:
                    self._update_progress(f"No search results found for question: {question}", 
                                        int(question_progress_base + 2),
                                        {"phase": "search_complete", "result_count": 0})
                    search_results = []  # Initialize to empty list instead of None
                    continue
                
                self._update_progress(f"Found {len(search_results)} results for question: {question}", 
                                    int(question_progress_base + 2),
                                    {"phase": "search_complete", "result_count": len(search_results)})
                
                logger.info(f"len search: {len(search_results)}")
                
                if len(search_results) == 0:
                    continue

                self._update_progress(f"Analyzing results for: {question}", 
                                     int(question_progress_base + 5),
                                     {"phase": "analysis"})


                try:
                    result = self.citation_handler.analyze_followup(
                        question, search_results, current_knowledge, nr_of_links=len(self.all_links_of_system)
                    )
                    links = extract_links_from_search_results(search_results)
                    self.all_links_of_system.extend(links)
                    section_links.extend(links)
                    formatted_links = ""  
                    if links:
                        formatted_links=format_links(links=links)
                    
                    logger.info(f"Generated questions: {formatted_links}")                           
                    if result is not None:
                        results_with_links = str(result["content"])
                        findings.append(
                            {
                                "phase": f"Follow-up {iteration}.{questions.index(question) + 1}",
                                "content": results_with_links,
                                "question": question,
                                "search_results": search_results,
                                "documents": result["documents"],
                            }
                        )

                        if settings.general.knowledge_accumulation != str(KnowledgeAccumulationApproach.NO_KNOWLEDGE.value):
                            current_knowledge = current_knowledge + "\n\n\n New: \n" + results_with_links
                        
                        if settings.general.knowledge_accumulation == str(KnowledgeAccumulationApproach.QUESTION.value):
                            logger.info("Compressing knowledge")
                            self._update_progress(f"Compress Knowledge for: {question}", 
                                        int(question_progress_base + 0),
                                        {"phase": "analysis"})
                            current_knowledge = self._compress_knowledge(current_knowledge , query, section_links)
                        
                        self._update_progress(f"Analysis complete for question: {question}", 
                                            int(question_progress_base + 10),
                                            {"phase": "analysis_complete"})
                except Exception as e:
                    error_msg = f"Error analyzing results: {str(e)}"
                    logger.info(f"ANALYSIS ERROR: {error_msg}")
                    self._update_progress(error_msg, 
                                        int(question_progress_base + 10),
                                        {"phase": "analysis_error", "error": str(e)})
            iteration += 1
            
            self._update_progress(f"Compressing knowledge after iteration {iteration}", 
                                 int((iteration / total_iterations) * 100 - 5),
                                 {"phase": "knowledge_compression"})
            logger.info(str(iteration))
            logger.info(settings.general.knowledge_accumulation)
            logger.info(str(KnowledgeAccumulationApproach.ITERATION.value))
            if settings.general.knowledge_accumulation == KnowledgeAccumulationApproach.ITERATION.value:
                try:
                    logger.info("ITERATION - Compressing Knowledge")
                    current_knowledge = self._compress_knowledge(current_knowledge , query, section_links)
                    logger.info("FINISHED ITERATION - Compressing Knowledge")
                except Exception as e:
                    error_msg = f"Error compressing knowledge: {str(e)}"
                    logger.info(f"COMPRESSION ERROR: {error_msg}")
                    self._update_progress(error_msg, 
                                        int((iteration / total_iterations) * 100 - 3),
                                        {"phase": "compression_error", "error": str(e)})


            
            self._update_progress(f"Iteration {iteration} complete", 
                                 int((iteration / total_iterations) * 100),
                                 {"phase": "iteration_complete", "iteration": iteration})
            
            try:
                formatted_findings = self._save_findings(findings, current_knowledge, query)
            except Exception as e:
                error_msg = f"Error saving findings: {str(e)}"
                logger.info(f"SAVE ERROR: {error_msg}")
                self._update_progress(error_msg, 
                                    int((iteration / total_iterations) * 100),
                                    {"phase": "save_error", "error": str(e)})
                formatted_findings = "Error: Could not format findings due to an error."

        self._update_progress("Research complete", 95, {"phase": "complete"})
        
        return {
            "findings": findings,
            "iterations": iteration,
            "questions": self.questions_by_iteration,
            "formatted_findings": formatted_findings,
            "current_knowledge": current_knowledge
        }

    def _save_findings(self, findings: List[Dict], current_knowledge: str, query: str):
        logger.info("Saving findings ...")
        self._update_progress("Saving research findings...", None)
        
        formatted_findings = format_findings_to_text(
            findings, current_knowledge, self.questions_by_iteration
        )
        safe_query = "".join(x for x in query if x.isalnum() or x in [" ", "-", "_"])[
            :50
        ]
        safe_query = safe_query.replace(" ", "_").lower()
        import local_deep_research.config as conf
        output_dir = f"{conf.get_config_dir()}/research_outputs"
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        filename = os.path.join(output_dir, f"formatted_output_{safe_query}.txt")

        with open(filename, "w", encoding="utf-8") as text_file:
            text_file.write(formatted_findings)
        logger.info("Saved findings")
        self._update_progress("Research findings saved", None, {"filename": filename})
        return formatted_findings