import re
from typing import Dict, List

from loguru import logger


LANGUAGE_CODE_MAP = {
    "english": "en",
    "french": "fr",
    "german": "de",
    "spanish": "es",
    "italian": "it",
    "japanese": "ja",
    "chinese": "zh",
    "hindi": "hi",
    "arabic": "ar",
    "bengali": "bn",
    "portuguese": "pt",
    "russian": "ru",
    "korean": "ko",
}


def remove_think_tags(text: str) -> str:
    # Remove paired <think>...</think> tags
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    # Remove any orphaned opening or closing think tags
    text = re.sub(r"</think>", "", text)
    text = re.sub(r"<think>", "", text)
    return text.strip()


def extract_links_from_search_results(search_results: List[Dict]) -> List[Dict]:
    """
    Extracts links and titles from a list of search result dictionaries.

    Each dictionary is expected to have at least the keys "title" and "link".

    Returns a list of dictionaries with 'title' and 'url' keys.
    """
    links = []
    if not search_results:
        return links

    for result in search_results:
        try:
            # Ensure we handle None values safely before calling strip()
            title = result.get("title", "")
            url = result.get("link", "")
            index = result.get("index", "")

            # Apply strip() only if the values are not None
            title = title.strip() if title is not None else ""
            url = url.strip() if url is not None else ""
            index = index.strip() if index is not None else ""

            if title and url:
                links.append({"title": title, "url": url, "index": index})
        except Exception:
            # Log the specific error for debugging
            logger.exception("Error extracting link from result")
            continue
    return links


def format_links_to_markdown(all_links: List[Dict]) -> str:
    formatted_text = ""
    logger.info(f"Formatting {len(all_links)} links to markdown...")

    if all_links:
        # Group links by URL and collect all their indices
        url_to_indices = {}
        for link in all_links:
            url = link.get("url")
            if url is None:
                url = link.get("link")
            index = link.get("index", "")
            # logger.info(f"URL \n {str(url)} ")
            if url:
                if url not in url_to_indices:
                    url_to_indices[url] = []
                url_to_indices[url].append(index)

        # Format each unique URL with all its indices
        seen_urls = set()  # Initialize the set here
        for link in all_links:
            url = link.get("url")
            if url is None:
                url = link.get("link")
            title = link.get("title", "Untitled")
            if url and url not in seen_urls:
                # Get all indices for this URL
                indices = sorted(
                    set(url_to_indices[url])
                )  # Sort for consistent ordering
                # Format as [1, 3, 5] if multiple indices, or just [1] if single
                indices_str = f"[{', '.join(map(str, indices))}]"
                # Add (source nr) as a visible fallback for humans
                formatted_text += f"{indices_str} {title} (source nr: {', '.join(map(str, indices))})\n   URL: {url}\n\n"
                seen_urls.add(url)

        formatted_text += "\n"

    return formatted_text


def format_findings(
    findings_list: List[Dict],
    synthesized_content: str,
    questions_by_iteration: Dict[int, List[str]],
) -> str:
    """Format findings into a detailed text output.

    Args:
        findings_list: List of finding dictionaries
        synthesized_content: The synthesized content from the LLM.
        questions_by_iteration: Dictionary mapping iteration numbers to lists of questions

    Returns:
        str: Formatted text output
    """
    logger.info(
        f"Inside format_findings utility. Findings count: {len(findings_list)}, Questions iterations: {len(questions_by_iteration)}"
    )
    formatted_text = ""

    # Extract all sources from findings
    all_links = []
    for finding in findings_list:
        search_results = finding.get("search_results", [])
        if search_results:
            try:
                links = extract_links_from_search_results(search_results)
                all_links.extend(links)
            except Exception:
                logger.exception("Error processing search results/links")

    # Start with the synthesized content (passed as synthesized_content)
    formatted_text += f"{synthesized_content}\n\n"

    # Add sources section after synthesized content if sources exist
    formatted_text += format_links_to_markdown(all_links)

    formatted_text += "\n\n"  # Separator after synthesized content

    # Add Search Questions by Iteration section
    if questions_by_iteration:
        formatted_text += "## SEARCH QUESTIONS BY ITERATION\n"
        formatted_text += "\n"
        for iter_num, questions in questions_by_iteration.items():
            formatted_text += f"\n #### Iteration {iter_num}:\n"
            for i, q in enumerate(questions, 1):
                formatted_text += f"{i}. {q}\n"
        formatted_text += "\n" + "\n\n"
    else:
        logger.warning("No questions by iteration found to format.")

    # Add Detailed Findings section
    if findings_list:
        formatted_text += "## DETAILED FINDINGS\n\n"
        logger.info(f"Formatting {len(findings_list)} detailed finding items.")

        for idx, finding in enumerate(findings_list):
            logger.debug(
                f"Formatting finding item {idx}. Keys: {list(finding.keys())}"
            )
            # Use .get() for safety
            phase = finding.get("phase", "Unknown Phase")
            content = finding.get("content", "No content available.")
            search_results = finding.get("search_results", [])

            # Phase header
            formatted_text += "\n"
            formatted_text += f"### {phase}\n"
            formatted_text += "\n\n"

            question_displayed = False
            # If this is a follow-up phase, try to show the corresponding question
            if isinstance(phase, str) and phase.startswith("Follow-up"):
                try:
                    parts = phase.replace("Follow-up Iteration ", "").split(".")
                    if len(parts) == 2:
                        iteration = int(parts[0])
                        question_index = int(parts[1]) - 1
                        if (
                            iteration in questions_by_iteration
                            and 0
                            <= question_index
                            < len(questions_by_iteration[iteration])
                        ):
                            formatted_text += f"#### {questions_by_iteration[iteration][question_index]}\n\n"
                            question_displayed = True
                        else:
                            logger.warning(
                                f"Could not find matching question for phase: {phase}"
                            )
                    else:
                        logger.warning(
                            f"Could not parse iteration/index from phase: {phase}"
                        )
                except ValueError:
                    logger.warning(
                        f"Could not parse iteration/index from phase: {phase}"
                    )
            # Handle Sub-query phases from IterDRAG strategy
            elif isinstance(phase, str) and phase.startswith("Sub-query"):
                try:
                    # Extract the index number from "Sub-query X"
                    query_index = int(phase.replace("Sub-query ", "")) - 1
                    # In IterDRAG, sub-queries are stored in iteration 0
                    if 0 in questions_by_iteration and query_index < len(
                        questions_by_iteration[0]
                    ):
                        formatted_text += (
                            f"#### {questions_by_iteration[0][query_index]}\n\n"
                        )
                        question_displayed = True
                    else:
                        logger.warning(
                            f"Could not find matching question for phase: {phase}"
                        )
                except ValueError:
                    logger.warning(
                        f"Could not parse question index from phase: {phase}"
                    )

            # If the question is in the finding itself, display it
            if (
                not question_displayed
                and "question" in finding
                and finding["question"]
            ):
                formatted_text += (
                    f"### SEARCH QUESTION:\n{finding['question']}\n\n"
                )

            # Content
            formatted_text += f"\n\n{content}\n\n"

            # Search results if they exist
            if search_results:
                try:
                    links = extract_links_from_search_results(search_results)
                    if links:
                        formatted_text += "### SOURCES USED IN THIS SECTION:\n"
                        formatted_text += (
                            format_links_to_markdown(links) + "\n\n"
                        )
                except Exception:
                    logger.exception(
                        f"Error processing search results/links for finding {idx}"
                    )
            else:
                logger.debug(f"No search_results found for finding item {idx}.")

            formatted_text += f"{'_' * 80}\n\n"
    else:
        logger.warning("No detailed findings found to format.")

    # Add summary of all sources at the end
    if all_links:
        formatted_text += "## ALL SOURCES:\n"
        formatted_text += format_links_to_markdown(all_links)
    else:
        logger.info("No unique sources found across all findings to list.")

    logger.info("Finished format_findings utility.")
    return formatted_text


def print_search_results(search_results):
    formatted_text = ""
    links = extract_links_from_search_results(search_results)
    if links:
        formatted_text = format_links_to_markdown(links)
    logger.info(formatted_text)
