#!/usr/bin/env python

import argparse
import logging
import os
import sys

from dotenv import load_dotenv
from llama_index.core import (
    SimpleDirectoryReader,
    StorageContext,
    VectorStoreIndex,
)
from llama_index.core.extractors import TitleExtractor
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.settings import Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.embeddings.ollama import OllamaEmbedding as _OllamaEmbedding


_UNICODE_REPLACEMENTS = str.maketrans(
    {
        "\u2013": "-",  # en-dash
        "\u2014": "--",  # em-dash
        "\u2026": "...",  # ellipsis
        "\u00a0": " ",  # non-breaking space
        "\u2018": "'",  # left single quote
        "\u2019": "'",  # right single quote
        "\u201c": '"',  # left double quote
        "\u201d": '"',  # right double quote
        # German umlauts: Ollama's mxbai BERT tokenizer has a bug where truncate=True
        # fails when umlauts appear alongside long dot sequences (e.g. TOC pages).
        "ä": "ae",
        "ö": "oe",
        "ü": "ue",
        "Ä": "Ae",
        "Ö": "Oe",
        "Ü": "Ue",
        "\u00df": "ss",  # ß
    }
)


class OllamaEmbedding(_OllamaEmbedding):
    """OllamaEmbedding with truncate=True and Unicode normalization.

    Ollama's mxbai BERT tokenizer crashes on some Unicode chars (e.g. en-dash)
    before truncation kicks in, so we replace them with ASCII equivalents first.
    """

    def get_general_text_embeddings(self, texts):
        clean = [t.translate(_UNICODE_REPLACEMENTS) for t in texts]
        result = self._client.embed(
            model=self.model_name,
            input=clean,
            truncate=True,
            keep_alive=self.keep_alive,
        )
        return [list(e) for e in result.embeddings]


from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.anthropic import Anthropic
from llama_index.llms.ollama import Ollama
from llama_index.llms.openai import OpenAI
from llama_index.readers.web import SimpleWebPageReader
from llama_index.vector_stores.milvus import MilvusVectorStore

load_dotenv()
parser = argparse.ArgumentParser(
    prog="vf-rag",
    description="Detailed information about viel-falter.at project",
)
parser.add_argument("-v", "--verbose", action="store_true")
parser.add_argument("-d", "--debug", action="store_true")
parser.add_argument("-r", "--regenerate", action="store_true")
parser.add_argument(
    "-l", "--llm", choices=["ollama", "openai", "claude"], default="claude"
)
parser.add_argument(
    "-e",
    "--embed",
    choices=["m3", "nomic", "mxbai", "base", "openai", "qwen3"],
    default="openai",
)
args = parser.parse_args()

if args.verbose or args.debug:
    logging.basicConfig(
        stream=sys.stdout, level=logging.DEBUG if args.debug else logging.INFO
    )
    logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

milvus_store = "http://localhost:19530"
ollama_host = os.environ.get("OLLAMA_HOST", "localhost:11434")
if not ollama_host.startswith("http"):
    ollama_host = f"http://{ollama_host}"
if ":" not in ollama_host.split("//")[-1]:
    ollama_host = f"{ollama_host}:11434"

dimensions = 768
if args.embed == "m3":
    Settings.embed_model = OllamaEmbedding(model_name="bge-m3", base_url=ollama_host)
    Settings.chunk_size = 512  # bge-m3 has 8192 token context
    Settings.chunk_overlap = 50
    dimensions = 1024
elif args.embed == "base":
    Settings.embed_model = HuggingFaceEmbedding(
        model_name="BAAI/bge-base-en-v1.5", cache_folder="/home/bernd/huggingface_cache"
    )
    Settings.chunk_size = (
        256  # BERT context is 512; tiktoken undercounts vs BERT tokenizer
    )
    Settings.chunk_overlap = 30
    dimensions = 768
elif args.embed == "qwen3":
    Settings.embed_model = OllamaEmbedding(
        model_name="qwen3-embedding", base_url=ollama_host
    )
    Settings.chunk_size = 512  # qwen3-embedding has large context
    Settings.chunk_overlap = 50
    dimensions = 4096
elif args.embed == "nomic":
    Settings.embed_model = OllamaEmbedding(
        model_name="nomic-embed-text", base_url=ollama_host
    )
    Settings.chunk_size = (
        256  # nomic BERT context is 512; tiktoken undercounts vs BERT tokenizer
    )
    Settings.chunk_overlap = 30
    dimensions = 768
elif args.embed == "mxbai":
    Settings.embed_model = OllamaEmbedding(
        model_name="mxbai-embed-large", base_url=ollama_host
    )
    Settings.chunk_size = (
        300  # mxbai BERT context is 512; tiktoken undercounts vs BERT tokenizer
    )
    Settings.chunk_overlap = 30
    dimensions = 1024
else:
    Settings.embed_model = OpenAIEmbedding()
    Settings.chunk_size = 1024
    Settings.chunk_overlap = 100
    dimensions = 1536


if args.llm == "ollama":
    Settings.llm = Ollama(model="qwen3:8b", request_timeout=360.0, base_url=ollama_host)
elif args.llm == "openai":
    Settings.llm = OpenAI(model="gpt-4o-mini")
else:
    Settings.llm = Anthropic(model="claude-sonnet-4-6")

logging.log(level=logging.INFO, msg=args)
logging.log(level=logging.DEBUG, msg=os.environ)
logging.log(level=logging.DEBUG, msg=Settings.llm)
logging.log(level=logging.DEBUG, msg=Settings.embed_model)

if args.regenerate:
    pdf_documents = SimpleDirectoryReader("pdfs").load_data(show_progress=True)
    webpages = [
        "https://viel-falter.at/",
        "https://viel-falter.at/monitoring/ziele-vision/",
        "https://viel-falter.at/projektpartner/",
        "https://viel-falter.at/news/",
        "https://viel-falter.at/tagfalter-exkursion/",
    ]
    web_documents = SimpleWebPageReader(html_to_text=True).load_data(webpages)
    documents = pdf_documents + web_documents

    vector_store = MilvusVectorStore(
        uri=milvus_store,
        dim=dimensions,
        overwrite=True,
        collection_name=args.embed,
        use_async_client=False,
    )
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    index = VectorStoreIndex.from_documents(
        documents=documents,
        transformations=[
            SentenceSplitter(
                chunk_size=Settings.chunk_size, chunk_overlap=Settings.chunk_overlap
            ),
            # TitleExtractor(),
            Settings.embed_model,
        ],
        storage_context=storage_context,
        show_progress=True,
    )
else:
    vector_store = MilvusVectorStore(
        uri=milvus_store,
        dim=dimensions,
        collection_name=args.embed,
        use_async_client=False,
    )
    index = VectorStoreIndex.from_vector_store(
        vector_store,
    )

query_engine = index.as_query_engine(similarity_top_k=6)


def ans(q):
    response = query_engine.query(q)
    print("\n#############################")
    print(f"Query: {q}")
    print("#############################")
    print(response)
    prompts_dict = query_engine.get_prompts()
    logging.log(level=logging.DEBUG, msg="=================== Prompts:")
    logging.log(level=logging.DEBUG, msg=prompts_dict)
    print("----------------------------")


# ans("Tell me about the goals and the project-partners of the viel-falter project")
# ans("Gib mir eine Liste aller Kooperationspartner")
# ans("What are latest news of the viel-falter project?")
# ans("What was the most found individual Tagfalterart in 2025 in the region Tyrol?")
# ans("What was the most commonly recorded butterfly species in Tyrol in 2025?")
ans("Welche Tagfalterart wurde im Jahr 2025 in Tirol am öftesten gefunden?")
# ans("Wer ist Peter Klimek?")
