#!/usr/bin/env python
import argparse
import logging
import sys

from dotenv import load_dotenv
from llama_index.core import (
    Settings,
    SQLDatabase,
)
from llama_index.core.query_engine import NLSQLTableQueryEngine

# from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.anthropic import Anthropic
from llama_index.llms.ollama import Ollama
from llama_index.llms.openai import OpenAI
from sqlalchemy import create_engine
from prompt_toolkit import PromptSession
from prompt_toolkit.styles import Style


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"],
    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))

engine = create_engine("sqlite:///data/vielfalter.db")

if args.llm == "ollama":
    Settings.llm = Ollama(
        model="llama3.1",
        request_timeout=360.0,
        base_url="http://localhost:11434",
    )
elif args.llm == "openai":
    Settings.llm = OpenAI(model="gpt-4o")
else:
    Settings.llm = Anthropic(model="claude-3-5-sonnet-20240620")

tables = ["surveys", "speciesList", "species", "sites"]
sql_database = SQLDatabase(engine, include_tables=tables)

query_engine = NLSQLTableQueryEngine(sql_database=sql_database, tables=tables)


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


session = PromptSession()
style = Style.from_dict({"": "#ff0066"})
try:
    while True:
        q = session.prompt("Query our viel-falter database: ", style=style)
        if q == "exit":
            break
        ans(q)
except (KeyboardInterrupt, EOFError):
    print("bye")

# ans("How many species were found in 2024?")
# ans("In which Nuts2 Region were most surveys conducted?")
# ans(
#     "How many species with german name 'Admiral' were found in surveys 2024 in nuts2 region of Vorarlberg and print the sql query?"
# )
