import os

import dspy
import turso

from config_loader import load_config
from embedding import LocalLMEmbeddings

CFG = load_config()

DATABASE_PATH = CFG["ingestion"]["db_path"]
DATABASE_NAME = CFG["ingestion"]["db_name"]
EMBEDDING_MODEL = CFG["models"]["embedding"]
API_BASE = CFG["api"]["base_url"]
RETRIEVAL_CONFIG = CFG["retrieval_agent"]
EXPANSION_CONFIG = CFG["expansion_agent"]


def retrieve_from_turso(embedded_question, k=5):
    query = f"""
    SELECT file_path, synopsis, tags, chunk_data,
    vector_distance_cos(embedding, vector32('{embedded_question}')) AS distance
    FROM notes
    ORDER BY distance ASC
    LIMIT {k};
    """
    con = turso.connect(DATABASE_PATH + DATABASE_NAME)
    cur = con.cursor()
    cur.execute(query)
    rows = cur.fetchall()
    return rows


class DnDContextQA(dspy.Signature):
    f"{RETRIEVAL_CONFIG['retrieval_signature']}"

    context = dspy.InputField(desc="Relevant chunks and metadata from the campaign notes.")
    question = dspy.InputField()
    answer = dspy.OutputField(desc="A detailed answer based on the notes, citing the source file.")


class ExpansionSignature(dspy.Signature):
    f"{EXPANSION_CONFIG['expansion_signature']}"
    question = dspy.InputField()
    answer = dspy.OutputField(
        desc="A list of questions that will be used to vector search the database."
    )


class DnDRAG(dspy.Module):
    def __init__(self):
        super().__init__()
        self.embeddings_model = LocalLMEmbeddings(
            model=EMBEDDING_MODEL,
            base_url=API_BASE,
            # batch_size=1,
        )
        self.retrieval_lm = dspy.LM(model=CFG["models"]["retrieval"], api_base=API_BASE)
        with dspy.context(lm=self.retrieval_lm, signature=ExpansionSignature):
            self.query_expander = dspy.Predict("question -> queries:list[str]")

        self.tools = [self.load_file]
        self.generate_answer = dspy.ReAct(signature=DnDContextQA, tools=self.tools)

    def forward(self, question):
        print("Enhancing Question")
        with dspy.context(lm=self.retrieval_lm):
            expanded_queries = self.query_expander(question=question).queries
        # print("Enhanced Queries:")
        # for q in expanded_queries:
        #     print("    ", q)
        all_embeddings = self.embeddings_model.embed_documents([question] + expanded_queries)
        # print(all_embeddings)
        all_results = []
        for embedded_question in all_embeddings:
            results = retrieve_from_turso(embedded_question, k=5)
            all_results.extend(results)

        seen = set()
        unique_results = []
        for row in all_results:
            key = (row[0], row[3])
            if key not in seen:
                seen.add(key)
                unique_results.append(row)

        context_parts = []
        for i, row in enumerate(unique_results):
            source = row[0]
            synopsis = row[1]
            tags = row[2]
            # entities = row[3]
            content = row[3]
            closeness = row[4]

            context_parts.append(f"""
--- Chunk {i + 1} from {source} ---
synopsis: {synopsis},
tags: {tags},
closeness: {closeness},
{content}
""")
        # entities: {entities},

        context = "\n\n".join(context_parts)

        prediction = self.generate_answer(context=context, question=question)
        return dspy.Prediction(answer=prediction.answer, context=context)

    def load_file(self, file_path) -> str | None:
        """Load and return specified file."""
        if os.path.exists(file_path):
            try:
                with open(file_path) as file:
                    return file.read()
            except Exception:
                return None
        else:
            return None