import json
import logging
import os
import re
from pathlib import Path
from collections import Counter

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer

from mm_agents.vlaa_gui.core.module import BaseModule
from mm_agents.vlaa_gui.memory.procedural_memory import PROCEDURAL_MEMORY
from mm_agents.vlaa_gui.utils.common_utils import (
    call_llm_safe,
    load_embeddings,
    load_knowledge_base,
    save_embeddings,
)

logger = logging.getLogger("desktopenv.agent")


class KnowledgeBase(BaseModule):
    def __init__(
        self,
        embedding_engine,
        local_kb_path: str,
        platform: str,
        engine_params: dict,
        save_knowledge: bool = True,
        lexical_weight: float = 0.0,
    ):
        super().__init__(engine_params, platform)

        self.local_kb_path = local_kb_path
        self.lexical_weight = lexical_weight

        # initialize embedding engine
        self.embedding_engine = embedding_engine

        # Initialize paths for different memory types
        self.episodic_memory_path = os.path.join(
            self.local_kb_path, self.platform, "episodic_memory.json"
        )
        self.narrative_memory_path = os.path.join(
            self.local_kb_path, self.platform, "narrative_memory.json"
        )
        self.embeddings_path = os.path.join(
            self.local_kb_path, self.platform, "embeddings.pkl"
        )

        # Initialize trajectory tracking
        self.task_trajectory = ""
        self.current_subtask_trajectory = ""
        self.current_search_query = ""

        self.rag_module_system_prompt = PROCEDURAL_MEMORY.RAG_AGENT.replace(
            "CURRENT_OS", self.platform
        )

        # All three agents share a generic RAG prompt that asks the agent to provide information for UI automation in CURRENT_OS
        self.query_formulator = self._create_agent(self.rag_module_system_prompt)
        self.llm_search_agent = self._create_agent(self.rag_module_system_prompt)
        self.knowledge_fusion_agent = self._create_agent(self.rag_module_system_prompt)

        self.save_knowledge = save_knowledge

        # Initialize TF-IDF vectorizer for lexical scoring
        self.tfidf_vectorizer = None
        self.lexical_corpus = []

    def _extract_keywords(self, text: str) -> list[str]:
        """Extract keywords from text using simple NLP preprocessing"""
        # Convert to lowercase and remove punctuation
        text = re.sub(r"[^\w\s]", " ", text.lower())
        # Split into words and filter out common stop words
        words = text.split()
        stop_words = {
            "the",
            "a",
            "an",
            "and",
            "or",
            "but",
            "in",
            "on",
            "at",
            "to",
            "for",
            "of",
            "with",
            "by",
            "is",
            "are",
            "was",
            "were",
            "be",
            "been",
            "being",
            "have",
            "has",
            "had",
            "do",
            "does",
            "did",
            "will",
            "would",
            "could",
            "should",
            "may",
            "might",
            "must",
            "can",
            "shall",
        }
        keywords = [word for word in words if len(word) > 2 and word not in stop_words]
        return keywords

    def _compute_lexical_scores(
        self, query: str, candidate_texts: list[str]
    ) -> np.ndarray:
        """Compute lexical similarity scores using TF-IDF and cosine similarity"""
        if not candidate_texts:
            return np.array([])

        try:
            # Extract keywords from query
            query_keywords = self._extract_keywords(query)
            if not query_keywords:
                return np.zeros(len(candidate_texts))

            # Extract keywords from candidate texts
            candidate_keywords = [
                self._extract_keywords(text) for text in candidate_texts
            ]

            # Create corpus for TF-IDF
            all_texts = [" ".join(query_keywords)] + [
                " ".join(keywords) for keywords in candidate_keywords
            ]

            # Fit TF-IDF vectorizer if not already done or if corpus changed
            if self.tfidf_vectorizer is None or self.lexical_corpus != all_texts:
                self.tfidf_vectorizer = TfidfVectorizer(
                    lowercase=True,
                    stop_words="english",
                    ngram_range=(1, 2),  # Include unigrams and bigrams
                    max_features=1000,
                )
                self.tfidf_matrix = self.tfidf_vectorizer.fit_transform(all_texts)
                self.lexical_corpus = all_texts

            # Compute cosine similarities between query and candidates
            query_vector = self.tfidf_matrix[0]  # First row is the query
            candidate_vectors = self.tfidf_matrix[1:]  # Rest are candidates

            similarities = cosine_similarity(query_vector, candidate_vectors).flatten()

            return similarities

        except Exception as e:
            logger.warning(f"Lexical scoring failed: {e}")
            return np.zeros(len(candidate_texts))

    def retrieve_knowledge(
        self, instruction: str, search_query: str, search_engine: str = "llm"
    ) -> tuple[str, str]:
        """Retrieve knowledge using search engine
        Args:
            instruction (str): task instruction
            observation (Dict): current observation
            search_engine (str): search engine to use"""

        # Use search engine to retrieve knowledge based on the formulated query
        search_results = self._search(instruction, search_query, search_engine)

        return search_query, search_results

    def formulate_query(self, instruction: str, observation: dict) -> str:
        """Formulate search query based on instruction and current state"""
        query_path = Path(self.local_kb_path) / self.platform / "formulate_query.json"
        try:
            with open(query_path, "r") as f:
                formulate_query = json.load(f)
        except Exception:
            formulate_query = {}

        if instruction in formulate_query:
            return formulate_query[instruction]

        self.query_formulator.reset()

        self.query_formulator.add_message(
            f"The task is: {instruction}\n"
            "To use google search to get some useful information, first carefully analyze "
            "the screenshot of the current desktop UI state, then given the task "
            "instruction, formulate a question that can be used to search on the Internet "
            "for information in helping with the task execution.\n"
            "The question should not be too general or too specific. Please ONLY provide "
            "the question.\nQuestion:",
            image_content=(
                observation["screenshot"] if "screenshot" in observation else None
            ),
            role="user",
        )

        search_query = self.query_formulator.get_response().strip().replace('"', "")
        print("search query: ", search_query)
        formulate_query[instruction] = search_query
        with open(query_path, "w") as f:
            json.dump(formulate_query, f, indent=2)

        return search_query

    def _search(self, instruction: str, search_query: str, search_engine: str) -> str:
        """Execute search using specified engine"""

        # Default to perplexica rag knowledge to see if the query exists
        file = (
            Path(self.local_kb_path)
            / self.platform
            / f"{search_engine}_rag_knowledge.json"
        )

        try:
            with open(file, "r") as f:
                exist_search_results = json.load(f)
        except Exception:
            exist_search_results = {}

        if instruction in exist_search_results:
            return exist_search_results[instruction]
        if search_engine.lower() == "llm":
            self.llm_search_agent.reset()
            # Use LLM's internal knowledge like a search engine
            self.llm_search_agent.add_message(search_query, role="user")
            search_results = self.llm_search_agent.get_response()
        else:
            raise ValueError(f"Unsupported search engine: {search_engine}")

        exist_search_results[instruction] = search_results.strip()
        with open(
            Path(self.local_kb_path)
            / self.platform
            / f"{search_engine}_rag_knowledge.json",
            "w",
        ) as f:
            json.dump(exist_search_results, f, indent=2)

        return search_results

    def retrieve_narrative_experience(self, instruction: str) -> tuple[str, str]:
        """Retrieve narrative experience using hybrid semantic + lexical scoring"""
        knowledge_base = load_knowledge_base(self.narrative_memory_path)

        if not knowledge_base:
            return "None", "None"

        embeddings = load_embeddings(self.embeddings_path)
        # Get or create instruction embedding
        instruction_embedding = embeddings.get(instruction)
        if instruction_embedding is None:
            instruction_embedding = self.embedding_engine.get_embeddings(instruction)
            embeddings[instruction] = instruction_embedding

        # Get or create embeddings for knowledge base entries
        candidate_embeddings = []
        keys = list(knowledge_base.keys())
        candidate_texts = []

        for key in keys:
            candidate_embedding = embeddings.get(key)
            if candidate_embedding is None:
                candidate_embedding = self.embedding_engine.get_embeddings(key)
                embeddings[key] = candidate_embedding
            candidate_embeddings.append(candidate_embedding)
            candidate_texts.append(knowledge_base[key])  # Store the actual content

        save_embeddings(self.embeddings_path, embeddings)

        # Compute semantic similarities (λ₁ = 1.0 - lexical_weight)
        semantic_similarities = cosine_similarity(
            instruction_embedding, np.vstack(candidate_embeddings)
        )[0]

        # Compute lexical similarities if lexical_weight > 0
        lexical_similarities = np.zeros(len(keys))
        if self.lexical_weight > 0:
            lexical_similarities = self._compute_lexical_scores(
                instruction, candidate_texts
            )

        # Hybrid scoring: S(q, m_k) = (1 - λ₂) · cos(e_q, v_k) + λ₂ · lexical_score
        hybrid_scores = (
            1 - self.lexical_weight
        ) * semantic_similarities + self.lexical_weight * lexical_similarities

        sorted_indices = np.argsort(hybrid_scores)[::-1]

        idx = 1 if keys[sorted_indices[0]] == instruction else 0
        return keys[sorted_indices[idx]], knowledge_base[keys[sorted_indices[idx]]]

    def retrieve_episodic_experience(self, instruction: str) -> tuple[str, str]:
        """Retrieve similar task experience using hybrid semantic + lexical scoring"""

        knowledge_base = load_knowledge_base(self.episodic_memory_path)
        if not knowledge_base:
            return "None", "None"

        embeddings = load_embeddings(self.embeddings_path)

        # Get or create instruction embedding
        instruction_embedding = embeddings.get(instruction)

        if instruction_embedding is None:
            instruction_embedding = self.embedding_engine.get_embeddings(instruction)
            embeddings[instruction] = instruction_embedding

        # Get or create embeddings for knowledge base entries
        candidate_embeddings = []
        keys = list(knowledge_base.keys())
        candidate_texts = []

        for key in keys:
            candidate_embedding = embeddings.get(key)
            if candidate_embedding is None:
                candidate_embedding = self.embedding_engine.get_embeddings(key)
                embeddings[key] = candidate_embedding
            candidate_embeddings.append(candidate_embedding)
            candidate_texts.append(knowledge_base[key])  # Store the actual content

        save_embeddings(self.embeddings_path, embeddings)

        # Compute semantic similarities (λ₁ = 1.0 - lexical_weight)
        semantic_similarities = cosine_similarity(
            instruction_embedding, np.vstack(candidate_embeddings)
        )[0]

        # Compute lexical similarities if lexical_weight > 0
        lexical_similarities = np.zeros(len(keys))
        if self.lexical_weight > 0:
            lexical_similarities = self._compute_lexical_scores(
                instruction, candidate_texts
            )

        # Hybrid scoring: S(q, m_k) = (1 - λ₂) · cos(e_q, v_k) + λ₂ · lexical_score
        hybrid_scores = (
            1 - self.lexical_weight
        ) * semantic_similarities + self.lexical_weight * lexical_similarities

        sorted_indices = np.argsort(hybrid_scores)[::-1]

        keys = list(knowledge_base.keys())
        idx = 1 if keys[sorted_indices[0]] == instruction else 0
        return keys[sorted_indices[idx]], knowledge_base[keys[sorted_indices[idx]]]

    def knowledge_fusion(
        self,
        observation: dict,
        instruction: str,
        web_knowledge: str,
        similar_task: str,
        experience: str,
    ) -> str:
        """Combine web knowledge with similar task experience"""

        self.knowledge_fusion_agent.reset()

        self.knowledge_fusion_agent.add_message(
            f"Task: {instruction}\n"
            f"**Web search result**:\n{web_knowledge}\n\n"
            f"**Retrieved similar task experience**:\n"
            f"Similar task:{similar_task}\n{experience}\n\n"
            f"Based on the web search result and the retrieved similar task experience, "
            f"if you think the similar task experience is indeed useful to the main task, "
            f"integrate it with the web search result. Provide the final knowledge in a numbered list.",
            image_content=(
                observation["screenshot"] if "screenshot" in observation else None
            ),
            role="user",
        )
        return self.knowledge_fusion_agent.get_response()

    def save_episodic_memory(self, subtask_key: str, subtask_traj: str) -> None:
        """Save episodic memory (subtask level knowledge).

        Args:
            subtask_key (str): Key identifying the subtask
            subtask_traj (str): Trajectory/experience of the subtask
        """
        if not self.save_knowledge:
            return

        try:
            kb = load_knowledge_base(self.episodic_memory_path)
        except Exception:
            kb = {}

        if subtask_key not in kb:
            subtask_summarization = self.summarize_episode(subtask_traj)
            kb[subtask_key] = subtask_summarization

            if self.save_knowledge:
                os.makedirs(os.path.dirname(self.episodic_memory_path), exist_ok=True)
                with open(self.episodic_memory_path, "w") as fout:
                    json.dump(kb, fout, indent=2)

        return kb.get(subtask_key)

    def save_narrative_memory(self, task_key: str, task_traj: str) -> None:
        """Save narrative memory (task level knowledge).

        Args:
            task_key (str): Key identifying the task
            task_traj (str): Full trajectory/experience of the task
        """
        if not self.save_knowledge:
            return

        try:
            kb = load_knowledge_base(self.narrative_memory_path)
        except Exception:
            kb = {}

        if task_key not in kb:
            task_summarization = self.summarize_narrative(task_traj)
            kb[task_key] = task_summarization

            if self.save_knowledge:
                os.makedirs(os.path.dirname(self.narrative_memory_path), exist_ok=True)
                with open(self.narrative_memory_path, "w") as fout:
                    json.dump(kb, fout, indent=2)

        return kb.get(task_key)

    def initialize_task_trajectory(self, instruction: str) -> None:
        """Initialize a new task trajectory.

        Args:
            instruction (str): The task instruction
        """
        self.task_trajectory = f"Task:\n{instruction}"
        self.current_search_query = ""
        self.current_subtask_trajectory = ""

    def update_task_trajectory(self, meta_data: dict) -> None:
        """Update the task trajectory with new metadata.

        Args:
            meta_data (Dict): Metadata from the agent's prediction
        """
        if not self.current_search_query and "search_query" in meta_data:
            self.current_search_query = meta_data["search_query"]

        self.task_trajectory += (
            "\n\nReflection:\n"
            + str(meta_data["reflection"])
            + "\n\n----------------------\n\nPlan:\n"
            + meta_data["executor_plan"]
        )

    def handle_subtask_trajectory(self, meta_data: dict) -> None:
        """Handle subtask trajectory updates based on subtask status.

        Args:
            meta_data (Dict): Metadata containing subtask information

        Returns:
            bool: Whether the subtask was completed
        """
        subtask_status = meta_data["subtask_status"]
        subtask = meta_data["subtask"]
        subtask_info = meta_data["subtask_info"]

        if subtask_status in ["Start", "Done"]:
            # If there's an existing subtask trajectory, finalize it
            if self.current_subtask_trajectory:
                self.current_subtask_trajectory += "\nSubtask Completed.\n"
                subtask_key = self.current_subtask_trajectory.split(
                    "\n----------------------\n\nPlan:\n"
                )[0]
                self.save_episodic_memory(subtask_key, self.current_subtask_trajectory)
                self.current_subtask_trajectory = ""
                return True

            # Start new subtask trajectory
            self.current_subtask_trajectory = (
                f"Task:\n{self.current_search_query}\n\n"
                f"Subtask: {subtask}\n"
                f"Subtask Instruction: {subtask_info}\n"
                f"----------------------\n\n"
                f"Plan:\n{meta_data['executor_plan']}\n"
            )
            return False

        elif subtask_status == "In":
            # Continue current subtask trajectory
            self.current_subtask_trajectory += (
                f"\n----------------------\n\nPlan:\n{meta_data['executor_plan']}\n"
            )
            return False

    def finalize_task(self) -> None:
        """Finalize the task by saving any remaining trajectories."""
        # Save any remaining subtask trajectory
        if self.current_subtask_trajectory:
            self.current_subtask_trajectory += "\nSubtask Completed.\n"
            subtask_key = self.current_subtask_trajectory.split(
                "\n----------------------\n\nPlan:\n"
            )[0]
            self.save_episodic_memory(subtask_key, self.current_subtask_trajectory)

        # Save the complete task trajectory
        if self.task_trajectory and self.current_search_query:
            self.save_narrative_memory(self.current_search_query, self.task_trajectory)

        # Reset trajectories
        self.task_trajectory = ""
        self.current_subtask_trajectory = ""
        self.current_search_query = ""

    def summarize_episode(self, trajectory):
        """Summarize the episode experience for lifelong learning reflection
        Args:
            trajectory: str: The episode experience to be summarized
        """

        # Create Reflection on whole trajectories for next round trial, keep earlier messages as exemplars
        self.episode_summarization_agent.add_message(trajectory)
        subtask_summarization = call_llm_safe(self.episode_summarization_agent)
        self.episode_summarization_agent.add_message(subtask_summarization)

        return subtask_summarization

    def summarize_narrative(self, trajectory):
        """Summarize the narrative experience for lifelong learning reflection
        Args:
            trajectory: str: The narrative experience to be summarized
        """
        # Create Reflection on whole trajectories for next round trial
        self.narrative_summarization_agent.add_message(trajectory)
        task_summarization = call_llm_safe(self.narrative_summarization_agent)

        return task_summarization
