import chainlit as cl
from chainlit.input_widget import Switch

import google.generativeai as gemini_client
from qdrant_client import QdrantClient

import textwrap
import os
import requests
import json

#API KEY
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
QDRANT_URL = os.environ.get("QDRANT_URL")
QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY")
VANNA_API_URL = os.environ.get("VANNA_API_URL")
VANNA_API_KEY = os.environ.get("VANNA_API_KEY")

#CONNECTION
client = QdrantClient(
    url=QDRANT_URL, 
    api_key=QDRANT_API_KEY,
)
gemini_client.configure(api_key=GEMINI_API_KEY)

def call_vanna_api(message):
    headers = {
        "Content-Type": "application/json",
        "VANNA-API-KEY": VANNA_API_KEY
    }
    data = {
        "message": message,
        "user_email": "rezky.yayang@gmail.com",  # Ganti dengan email pengguna
        "agent_id": "landw",
        "acceptable_responses": ["text"],
    }

    try:
        response = requests.post(VANNA_API_URL, headers=headers, data=json.dumps(data), stream=True)
        response.raise_for_status()

        results = []
        for line in response.iter_lines():
            if line:
                decoded_line = line.decode('utf-8')
                if decoded_line.startswith("data:"):
                    data_string = decoded_line[5:].strip()
                    try:
                        data = json.loads(data_string)
                        if data['type'] == 'text':
                            results.append(data['text'])  # Append the text response to the results list
                        elif data['type'] == 'end':
                            # Return only the last 'text' type response
                            return results[-1] if results else "No valid text response found"
                        elif data['type'] == 'error':
                            results.append(f"**Error:**\n{data['error']}")
                    except json.JSONDecodeError as e:
                        results.append(f"Error decoding JSON: {e} - Original data: {data_string}")
        return "\n\n".join(results)
    except requests.exceptions.RequestException as e:
        return f"An error occurred: {e}"
    except Exception as e:
        return f"An unexpected error occurred: {e}"


@cl.set_chat_profiles
async def chat_profile():
    return [
        cl.ChatProfile(
            name="Databot LAN",
            markdown_description="Temukan data statistik sektoral LANRI dengan bantuan AI",
            icon="/public/terminal.svg",
            starters=[
                cl.Starter(
                    label="JF AK menurut jenis kelamin?",
                    message="jumlah jabatan fungsional analis kebijakan menurut jenis kelamin?",
                    icon="/public/idea.svg",
                ),
                cl.Starter(
                    label="Alumni Latsar CPNS tahun 2023?",
                    message="jumlah alumni Pelatihan Dasar CPNS tahun 2023 menurut asal instansi?",
                    icon="/public/learn.svg",
                ),
                cl.Starter(
                    label="PPPK mengakses Swajar Orientasi?",
                    message="jumlah PPPK yang mengakses swajar orientasi PPPK berdasarkan nama jabatan?",
                    icon="/public/terminal.svg",
                    ),
                cl.Starter(
                    label="Alumni PKN I per tahun?",
                    message="jumlah alumni PKN I berdasarkan tahun diklat?",
                    icon="/public/write.svg",
                    )
                    ],
        ), 
        cl.ChatProfile(
            name="Infobot LAN",
            markdown_description="Cari tahu informasi umum mengenai LAN RI dengan bantuan AI",
            icon="/public/write.svg",
            starters=[
                cl.Starter(
                    label="Apa itu Bangkom ASN?",
                    message="Apa yang dimaksud dengan Pengembangan Kompetensi?",
                    icon="/public/terminal.svg",
                    ),
                cl.Starter(
                    label="Minimal JP Bangkom ASN?",
                    message="Berapa minimal jam pelajaran pengembangan kompetensi bagi PNS dalam setahun?",
                    icon="/public/idea.svg",
                ),
                cl.Starter(
                    label="Masa prajabatan CPNS?",
                    message="Berapa lama masa prajabatan bagi seorang CPNS?",
                    icon="/public/learn.svg",
                ),
                cl.Starter(
                    label="Berapa jam pelajaran PKN I?",
                    message="Berapa jam pelajaran yang didapatkan ketika seorang ASN mengikuti pelatihan kepemimpinan nasional tingkat 1?",
                    icon="/public/write.svg",
                    )
                    ],
        ),
    ]

@cl.on_message
async def main(message: cl.Message):
    if cl.user_session.get("chat_profile") == "Databot LAN":
        human_query = message.content

        # Panggil API Vanna AI untuk mendapatkan respons
        vanna_response = call_vanna_api(human_query)

        await cl.Message(content=vanna_response).send()

    elif cl.user_session.get("chat_profile") == "Infobot LAN":
        ask = message.content

        search_result = client.search(
                collection_name='lan_knowledge',
                query_vector=gemini_client.embed_content(
                    model="models/embedding-001",
                    content=ask,
                    task_type="retrieval_query",
                )["embedding"],
                limit = 2,
            )

        def make_prompt(question, relevant_passage,data_source):
          escaped = relevant_passage.replace("'", "").replace('"', "").replace("\n", " ")
          prompt = textwrap.dedent("""You are a helpful and informative bot that answers questions using text from the reference passage included below. \
          Be sure to respond in a complete sentence, being comprehensive, including all relevant background information. \
          However, you are talking to a non-technical audience, so be sure to break down complicated concepts and \
          strike a friendly and converstional tone. \
          If the passage is irrelevant to the answer, you may cut or ignore it. \
          Include the data source at the end of the answer, separate by enter \
          QUESTION: '{question}'
          PASSAGE: '{relevant_passage}'
          Sumber Data: {data_source}

            ANSWER: (answer with Bahasa Indonesia)
            if the passage irrelevant, answer: Maaf, sementara Infobot LANRI belum dapat menjawab pertanyaan tersebut. Silakan mengajukan pertanyaan lainnya.
          """).format(question=question, relevant_passage=escaped,data_source=data_source)

          return prompt

        prompt = make_prompt(question=ask,
                         relevant_passage = search_result[0].payload['answer'],
                         data_source = search_result[0].payload['pasal'] + '  ' + search_result[0].payload['peraturan'])
        model = gemini_client.GenerativeModel('models/gemini-2.0-flash')
        answer = model.generate_content(prompt)

        await cl.Message(
            content=f"{answer.text}").send()