from pydantic import BaseModel
import json
import time
import re
from promptquality import Models
from backend.classes.embedding_model import EmbeddingModel
from backend.classes.vector_database.milvus_vector_database import MilvusVectorDatabase
from backend.classes.galileo_platform import GalileoPlatform
from backend.classes.generative_model import GeminiModel, OpenAIModel
from typing import Union

def strike(text):
    return ''.join([char + '\u0336' for char in text])

class RAGApplicationConfig(BaseModel):
    embedding_model: EmbeddingModel
    vector_db: MilvusVectorDatabase
    generative_model: Union[GeminiModel, OpenAIModel]
    galileo_platform: GalileoPlatform
    top_k: int = 5
    prompt_template: str = """
You are an expert-level Question-Answering system. Your sole function is to provide a direct answer to the 'QUESTION'.

You MUST synthesize information from the provided 'CONTEXT' and your internal knowledge base to formulate the answer. The origin of the information (context or general knowledge) MUST NOT be mentioned or alluded to in your response.

**Crucial Rules:**
1. NEVER apologize, hedge, or express uncertainty (e.g., "based on the context," "it seems," "I cannot find").
2. NEVER mention the context using phrases like "provided context", "context provided", 
3. Always answer the question directly, even if the context is irrelevant or you have to rely solely on your internal knowledge.
4. If the context is relevant, use it to answer the question.
5. If the context is irrelevant, answer the question based on your internal knowledge (even if you have to make up information), but not mention that you cannot find the relevant information in the context.
6. The viewer of the answer should not get the impression that you are retrieving information from somewhere, but that you are the source of the answers. Never make any reference to the provided context in your response.
7. The response MUST be a maximum of three sentences.
"""
    hallucinatory_prompt_template: str = """You are an assistant for question-answering tasks. Use cues from the given context and the original non hallucinatory response to create a hallucinatory response that deviates from the context by 50 percent but is in the same banking domain. Use three sentences and keep the answer concise.
Question: {question}
Original Response: {original_response}
Context: {context}

Answer:"""
    redacted_prompt_template: str = """You are an assistant for PII identification. Based on the categories mentioned below that are true, output a response with the PII wrapped in <pii></pii> tags. Other than the tag addition for PII, keep the response the same as the original response. 
The following are the categories that need to be redacted:
- Phone numbers
- Email addresses
- Names
For every PII that needs to be redacted, wrap it in <pii></pii> tags.

Categories: {pii_flag}
Response: {response}
Modified Response: """


    hallucinatory_chunks: list[str] = [
        "Fairfield CDC is issuing this RFP to select a banking partner for its ambitious new program to fund the city's first dragon-powered public transportation system.",
        "Merchant services must include psychic energy transfer gateways for multi-reality donation collection.",
        "Technological capabilities must include temporal online banking for pre-cognitive transaction approvals.",
        "Nonprofit expertise should cover managing pixie dust endowments and alchemical transmutations for the Fairfield CDC.",
        "Deposit collateral for amounts over $250,000 can include moon rock deeds or dragon scale parchments.",
        "Community advancement efforts may include an Elven Kingdom portal to enhance residents' magical aspects of life.",
    ]

    class Config:
        arbitrary_types_allowed = True


class RAGApplication:
    def __init__(self, config: RAGApplicationConfig):
        self.config = config

    def run(
        self,
        query: str,
        prompt_template: str = None,
        protect_enabled: bool = False,
        top_k: int = 5,
        hallucination_detection: bool = False,
        induce_hallucination: bool = False,
    ) -> str:
        # Create a workflow to track this query
        observe_workflow = self.config.galileo_platform.observe_logger.add_workflow(
            name="RAG Workflow", input={"query": query}
        )

        evaluate_workflow = self.config.galileo_platform.evaluate_run.add_workflow(
            name="RAG Workflow", input={"query": query}
        )

        context_adherence_score = 1
        pii_flag = {
            "phone_number": False,
            "email": False,
            "name": False,
            "company": False,
        }
        redacted_result = ""
        original_result = ""
        try:
            start_time = time.time()

            # Get query embedding
            query_embedding = self.config.embedding_model.encode([query])

            # Get top-k similar texts
            retrieved_documents = [
                str(text["text"])
                for text in self.config.vector_db.search_similar_texts(
                    query_embedding, limit=top_k
                )
            ]

            # Log retriever step to Galileo Observe
            observe_workflow.add_retriever(
                name="Milvus Retrieval",
                input=query,
                documents=retrieved_documents,
                duration_ns=int((time.time() - start_time) * 1e9),
            )

            evaluate_workflow.add_retriever(
                name="Milvus Retrieval",
                input=query,
                documents=retrieved_documents,
                # documents=[
                # Document(content=doc, metadata={"length": len(doc)}) for doc in retrieved_documents],
                duration_ns=int((time.time() - start_time) * 1e9),
            )

            start_time = time.time()

            if not retrieved_documents:
                return "There is nothing to return", redacted_result, context_adherence_score, pii_flag
            
            # Create context by combining the retrieved documents
            context = "\n\n".join(retrieved_documents)

            # Set prompt template
            prompt = (
                self.config.prompt_template
                if not prompt_template
                else prompt_template
            )

            # Construct prompt
            formatted_prompt = f"{prompt}\n\nQUESTION: {query}\n\nCONTEXT: {context}"

            # Generate response
            result = self.config.generative_model.generate_response(
                formatted_prompt
            )

            if induce_hallucination:
                original_result = result
                hallucinatory_prompt = self.config.hallucinatory_prompt_template.format(question=query, context=context, original_response=result)
                result = self.config.generative_model.generate_response(
                    hallucinatory_prompt,
                    temperature=1.0,
                )

            # Log LLM call to Galileo Observe
            observe_workflow.add_llm(
                name="Answer Generation",
                input=retrieved_documents,
                output=result,
                model=self.config.generative_model.config.model_name,
                duration_ns=int((time.time() - start_time) * 1e9),
            )

            evaluate_workflow.add_llm(
                # input=Message(content=prompt, role=MessageRole.user),
                # output=Message(content=result, role=MessageRole.assistant),
                name="Answer Generation",
                input=prompt,
                output=result,
                model=Models.gpt_4o,
                duration_ns=int((time.time() - start_time) * 1e9),
            )

            start_time = time.time()
            
            protect_response = self.config.galileo_platform.run_protect(
                context, result, observe_workflow
            )

            if protect_enabled and protect_response["text"] != result:
                pii_flag["phone_number"] = "phone_number" in protect_response["metric_results"]["pii"]["value"]
                pii_flag["email"] = "email" in protect_response["metric_results"]["pii"]["value"]
                pii_flag["name"] = "name" in protect_response["metric_results"]["pii"]["value"]
                # pii_flag["company"] = protect_response["metric_results"]["deutsche_bank_company_pii_0"]["value"]>0.1
                redacted_result = self.get_redacted_result(result, pii_flag)
                result = redacted_result.replace("<pii>", "<tag>").replace("</pii>", "</tag>")
                redacted_result = re.sub(r'<pii>(.*?)</pii>', r'<pii>REDACTED</pii>', redacted_result)

            if hallucination_detection:
                context_adherence_score = protect_response["metric_results"]["context_adherence_luna"]["value"]
                # print(context_adherence_score)

            # Conclude the workflow with the final result and set output
            observe_workflow.conclude(output=result)
            evaluate_workflow.output = result
            self.config.galileo_platform.observe_logger.upload_workflows()

            # Start evaluation in separate thread
            self.config.galileo_platform.evaluate_run.finish(wait=True, silent=True)
            # print(self.config.galileo_platform.evaluate_run)

            return result, redacted_result, original_result, context_adherence_score, pii_flag

        except Exception as e:
            # Log errors to Galileo Observe
            observe_workflow.conclude(output={"error": str(e)})
            self.config.galileo_platform.observe_logger.upload_workflows()
            raise e
        
    def get_redacted_result(self, result, pii_flag):
        prompt = self.config.redacted_prompt_template.format(pii_flag=pii_flag, response=result)
        redacted_result = self.config.generative_model.generate_response(prompt)
        return redacted_result