import os
import logging
import requests
from typing import Dict, List, Any, Optional

# haystack/ neo4j integration components
from neo4j_haystack import Neo4jDocumentStore, Neo4jEmbeddingRetriever

from haystack import component, Pipeline
from haystack.dataclasses import Document, ChatMessage, TextContent
from haystack.components.agents import Agent
from haystack.components.builders import PromptBuilder
from haystack.components.converters import HTMLToDocument
from haystack.components.rankers import TransformersSimilarityRanker
from haystack.components.fetchers import LinkContentFetcher
from haystack.components.preprocessors import DocumentCleaner
from haystack.components.rankers import TransformersSimilarityRanker
from haystack.components.websearch import SerperDevWebSearch
from haystack.components.evaluators import (
    FaithfulnessEvaluator,
    ContextRelevanceEvaluator,
)
from haystack.core import SuperComponent
from haystack.dataclasses import ChatMessage, ToolCall, ToolCallResult, TextContent
from haystack.tools.component_tool import ComponentTool
from haystack.utils import Secret
from haystack_integrations.components.embedders.mistral.document_embedder import (
    MistralDocumentEmbedder,
)
from haystack_integrations.components.embedders.mistral.text_embedder import (
    MistralTextEmbedder,
)
from haystack.components.generators.utils import print_streaming_chunk
from haystack.dataclasses import ChatMessage
from haystack_integrations.components.generators.mistral import MistralChatGenerator
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack.components.builders import ChatPromptBuilder

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------
# Document Store
# ---------------------------------------------------------------------


NEO4J_URI = os.environ["NEO4J_URI"]

NEO4J_USER = os.environ["NEO4J_USERNAME"]
NEO4J_PASS = os.environ["NEO4J_PASSWORD"]
SERPERDEV_API_KEY = os.environ["SERPERDEV_API_KEY"]

document_store = Neo4jDocumentStore(
    url=NEO4J_URI,
    username=NEO4J_USER,
    password=NEO4J_PASS,
    database="neo4j",
    index="document-embeddings",
    embedding_field="embedding",
    embedding_dim=1024,
    node_label="Document",
)


# ---------------------------------------------------------------------
# Helper Components
# ---------------------------------------------------------------------
@component
class DocumentPassthrough:
    @component.output_types(documents=List[Document])
    def run(self, documents: List[Document]) -> Dict[str, List[Document]]:
        return {"documents": documents}


@component
class TemplateVariablesProvider:
    @component.output_types(template_variables=Optional[Dict[str, Any]])
    def run(self, query: str, documents: List[Document]) -> Dict[str, Any]:
        return {"template_variables": {"query": query, "documents": documents}}


@component
class HFAPIRanker:
    """
    Ranks Documents using Hugging Face Inference API
    (e.g. sentence-transformers/all-MiniLM-L6-v2)
    """

    def __init__(
        self,
        model_id: str = "sentence-transformers/all-MiniLM-L6-v2",
        token: Optional[str] = None,
        top_k: int = 5,
        timeout: int = 15,
    ):
        self.url = f"https://api-inference.huggingface.co/models/{model_id}"
        self.headers = {"Authorization": f"Bearer {token or os.getenv('HF_TOKEN')}"}
        self.top_k = top_k
        self.timeout = timeout

    @component.output_types(documents=List[Document])
    def run(self, query: str, documents: List[Document]) -> Dict[str, List[Document]]:
        if not documents:
            return {"documents": []}

        payload = {
            "inputs": {
                "source_sentence": query,
                "sentences": [d.content for d in documents],
            }
        }

        resp = requests.post(
            self.url, headers=self.headers, json=payload, timeout=self.timeout
        )
        resp.raise_for_status()
        scores = resp.json()  # list of floats

        for d, s in zip(documents, scores):
            d.score = float(s)

        ranked = sorted(documents, key=lambda d: d.score, reverse=True)[: self.top_k]
        return {"documents": ranked}


# ---------------------------------------------------------------------
# Globals
# ---------------------------------------------------------------------
_rag_initialized: bool = False
_agent_instance: Optional[Agent] = None


# ---------------------------------------------------------------------
# Initialization
# ---------------------------------------------------------------------
def initialize_rag_components(mistral_api_key: str):
    """Initialize embedding pipeline, web pipeline, and agent."""
    global _rag_initialized, _agent_instance
    if _rag_initialized:
        logger.info("RAG components already initialized.")
        return

    os.environ["MISTRAL_API_KEY"] = mistral_api_key
    hf_token = os.getenv("HF_TOKEN")

    # ---------------------- Embedding pipeline ----------------------
    embedder_emb = MistralTextEmbedder()
    retriever_emb = Neo4jEmbeddingRetriever(document_store=document_store)
    # ranker_emb = HFAPIRanker(token=hf_token, top_k=2)
    ranker_emb = TransformersSimilarityRanker(
        model="intfloat/simlm-msmarco-reranker", top_k=3
    )

    prompt_template_emb = [
        ChatMessage.from_user(
            """
        You are an AI Assistant with access to official documents about the Government of Uganda's services, programs, policies, and initiatives, especially those designed to help citizens interact with government effectively through the Government Citizen Interaction Centre (GCIC).
        Your task is to answer user questions strictly based on the documents provided below, offering practical guidance, step-by-step instructions, and clear advice that helps citizens access and use government services.
        
        Guidelines:
        
        Use only the content from the provided Documents.
        
        Do NOT rely on prior or external knowledge.
        
        Do NOT ask the user for additional information.
        
        Include inline HTML links for referencing URL sources in the answer, using the URLs provided in the Documents. Use the document's title as the anchor text. If the title is missing, use the domain name of the document's URL as the anchor text.
        
        Each fact you refer to should be followed by the corresponding reference.
        
        Output the answer in a structured markdown format.
        
        Use bullet lists, numbered steps, or tables whenever it makes the guidance clearer.
        
        Focus on actionable advice for citizens, such as:
        
        How to access services or programs
        
        What steps to follow to engage with government offices
        
        Important deadlines, requirements, or tips for using services efficiently
        
        Do not add a references section at the end of the answer, just use references within the body of text.
        
        If a definitive answer cannot be found in the Documents, respond with:
        Final Answer: inconclusive
        
        Always end your answer with this disclaimer:
        Disclaimer: This is AI generated content — please use it with caution.
        
        Documents:
        {% for doc in documents %}
        Source: <a href="{{ doc.meta.source_url }}"></a><br>
        Title: <a href="{{ doc.meta.title }}"></a><br>
        Date: <a href="{{ doc.meta.date }}"></a><br>
        
        {{ doc.content }}
        {% endfor %}
        
        Question: {{ query }}
        
        Answer:
        """
        )
    ]
    prompt_builder_emb = ChatPromptBuilder(template=prompt_template_emb)
    generator_emb = MistralChatGenerator(model="mistral-small-latest")
    emb_template_provider = TemplateVariablesProvider()

    emb_pipeline = Pipeline()
    emb_pipeline.add_component("embedder", embedder_emb)
    emb_pipeline.add_component("retriever", retriever_emb)
    emb_pipeline.add_component("reranker", ranker_emb)
    emb_pipeline.add_component("output_docs", DocumentPassthrough())
    emb_pipeline.add_component("template_provider", emb_template_provider)
    emb_pipeline.add_component("prompt_builder", prompt_builder_emb)
    emb_pipeline.add_component("generator", generator_emb)

    emb_pipeline.connect("embedder.embedding", "retriever.query_embedding")
    emb_pipeline.connect("retriever.documents", "reranker.documents")
    emb_pipeline.connect("reranker.documents", "output_docs.documents")
    emb_pipeline.connect("reranker.documents", "template_provider.documents")
    emb_pipeline.connect(
        "template_provider.template_variables", "prompt_builder.template_variables"
    )
    emb_pipeline.connect("prompt_builder.prompt", "generator.messages")
    emb_pipeline.warm_up()

    # ---------------------- Web pipeline ----------------------
    web_search = SerperDevWebSearch(top_k=2)
    fetcher = LinkContentFetcher()
    converter = HTMLToDocument()
    # ranker_web = HFAPIRanker(token=hf_token, top_k=2)
    ranker_web = TransformersSimilarityRanker(
        model="intfloat/simlm-msmarco-reranker", top_k=3
    )

    prompt_template_web = [
        ChatMessage.from_user(
            """
    You are a helpful AI assistant focused on guiding citizens in accessing and using Government of Uganda services, programs, and initiatives. Use only the web search documents below to answer the user's question.
    
    Instructions:
    
    Use only the content from the provided documents.
    
    Do NOT use prior knowledge or make assumptions.
    
    Provide practical, step-by-step guidance that helps citizens:
    
    Access services or programs
    
    Complete forms or applications
    
    Understand requirements, deadlines, or procedures
    
    Make the most of available government initiatives
    
    When referencing a source, include an inline HTML link using the document's URL. Each fact you refer to should be followed by the corresponding reference.
    
    If a clear title can be inferred from the document content or URL, use it as the anchor text. If no clear title can be inferred, use the domain name of the URL as the anchor text.
    
    Output the answer in a structured markdown format.
    
    Use bullet lists, numbered steps, or tables whenever it makes the guidance clearer.
    
    Avoid technical jargon; write in a friendly, easy-to-follow tone.
    
    Do not add a references section at the end of the answer, just use references within the body of text.
    
    If no relevant information is found in the documents, respond with:
    Final Answer: inconclusive
    
    Always end your answer with this disclaimer:
    Disclaimer: This is AI generated content — please use it with caution.
    
    Documents:
    {% for doc in documents %}
    <b>Source:</b> <a href="{{ doc.meta.url }}">{{ doc.meta.url }}</a><br>
    <p>{{ doc.content }}</p><br>
    {% endfor %}
    
    Question: {{ query }}
    
    Answer:
    """
        )
    ]
    prompt_builder_web = ChatPromptBuilder(template=prompt_template_web)
    web_template_provider = TemplateVariablesProvider()
    generator_web = MistralChatGenerator(model="mistral-small-latest")

    web_pipeline = Pipeline()
    web_pipeline.add_component("search", web_search)
    web_pipeline.add_component("fetcher", fetcher)
    web_pipeline.add_component("converter", converter)
    web_pipeline.add_component("ranker", ranker_web)
    web_pipeline.add_component("output_docs", DocumentPassthrough())
    web_pipeline.add_component("template_provider", web_template_provider)
    web_pipeline.add_component("prompt_builder", prompt_builder_web)
    web_pipeline.add_component("generator", generator_web)

    web_pipeline.connect("search.links", "fetcher.urls")
    web_pipeline.connect("fetcher.streams", "converter.sources")
    web_pipeline.connect("converter.documents", "ranker.documents")
    web_pipeline.connect("ranker.documents", "output_docs.documents")
    web_pipeline.connect("ranker.documents", "template_provider.documents")
    web_pipeline.connect(
        "template_provider.template_variables", "prompt_builder.template_variables"
    )
    web_pipeline.connect("prompt_builder.prompt", "generator.messages")
    web_pipeline.warm_up()

    # ---------------------- Agent ----------------------
    chat_generator = MistralChatGenerator(model="mistral-small-latest")

    embedding_super = SuperComponent(
        pipeline=emb_pipeline,
        input_mapping={
            "query": ["embedder.text", "template_provider.query", "reranker.query"]
        },
        output_mapping={"generator.replies": "replies"},
    )

    web_super = SuperComponent(
        pipeline=web_pipeline,
        input_mapping={
            "query": ["search.query", "ranker.query", "template_provider.query"]
        },
        output_mapping={"generator.replies": "replies"},
    )

    embedding_tool = ComponentTool(
        name="embedding_search",
        component=embedding_super,
        description=(
            "Answer questions using information retrieved from an internal document store containing content "
            "about the Government of Uganda’s services, programs, policies, and initiatives, especially those relevant "
            "to improving citizen interactions through the Government Citizen Interaction Centre (GCIC). Answers are strictly "
            "based on retrieved documents using semantic similarity, with practical guidance for citizens. References are included as HTML links."
        ),
    )

    web_tool = ComponentTool(
        name="web_search",
        component=web_super,
        description=(
            "Retrieve potentially relevant information from the web about the Government of Uganda’s services, programs, and initiatives, "
            "with emphasis on citizen engagement and access to government services. Results are based on live internet content and may include a variety of sources. "
            "Answers include actionable guidance for citizens where possible. References are provided as HTML links using either inferred titles or domain names."
        ),
    )

    system_prompt = """
    You are an assistant with two tools:
    - embedding_search: retrieves answers from internal GCIC documents.
    - web_search: retrieves live information from the web.
    
    Rules:
    - Always use embedding_search first.
    - If embedding_search returns relevant results, answer immediately using only those results (do not call web_search).
    - If embedding_search has no results, then use web_search.
    - If neither tool provides an answer, respond with:
    Final Answer: inconclusive
    
    Answer formatting:
    - Always use **Markdown** to structure responses.
    - Use headings (##) for sections when explaining.
    - Use bullet points or numbered lists for clarity.
    - Use bold text for important terms or highlights.
    - Use markdown tables for structured information (e.g., fees, service times).
    - Use markdown links `[link text](url)` for references.
    - Keep answers conversational, clear, and easy to read — avoid technical jargon.

    
    Always end every answer with:
    Disclaimer: This is AI generated content — please use it with caution.
    """

    agent_instance = Agent(
        chat_generator=chat_generator,
        tools=[embedding_tool],
        system_prompt=system_prompt,
    )
    agent_instance.warm_up()

    _agent_instance = agent_instance
    _rag_initialized = True
    logger.info("RAG agent initialized successfully.")


# ---------------------------------------------------------------------
# Accessors
# ---------------------------------------------------------------------
def get_rag_components():
    if not _rag_initialized:
        return {"llm": None, "agent": None}
    return {"llm": _agent_instance, "agent": _agent_instance}


def get_rag_response_dynamic(
    user_query_english: str, chat_history_english: list, on_token=None
):
    if not _rag_initialized or _agent_instance is None:
        raise RuntimeError("RAG components are not initialized.")

    # Convert chat history (list of dict or str) to ChatMessage
    chat_msgs = []
    for item in chat_history_english:
        if isinstance(item, ChatMessage):
            chat_msgs.append(item)
        elif isinstance(item, dict):
            chat_msgs.append(
                ChatMessage(role=item.get("role"), content=item.get("content"))
            )
        elif isinstance(item, str):
            chat_msgs.append(ChatMessage.from_user(item))

    # Append current user query
    chat_msgs.append(ChatMessage.from_user(user_query_english))

    def my_streaming_callback(chunk):
        # Process each chunk as it arrives
        if chunk.content:
            print(chunk.content, end="", flush=True)
            if on_token:
                # Add this line
                on_token(chunk.content)  # Add this line

    result = _agent_instance.run(
        messages=chat_msgs,
        max_steps=6,
        streaming_callback=my_streaming_callback,
    )

    final_answer = None
    for msg in result.get("messages", []):
        if msg.role == "assistant" and msg.text:
            final_answer = msg.text
            break

    return final_answer or "Final Answer: inconclusive"
