import gradio as gr
from pathlib import Path
from sambanova import SambaNova
from langchain_huggingface import HuggingFaceEmbeddings

from chatbot import (
    load_config,
    build_rag_corpus,
    retrieve_relevant_chunks,
    build_prompt,
    ask_model,
    format_answer,
)

CONFIG_PATH = Path(__file__).parent / "config.yaml"
RESOURCE_STATE = {}


def init_resources():
    if RESOURCE_STATE:
        return RESOURCE_STATE

    if not CONFIG_PATH.exists():
        raise FileNotFoundError(f"Missing config file: {CONFIG_PATH}")

    config = load_config(CONFIG_PATH)
    llm_api_key = config.get("sambanova_api_key")
    website = config.get("website")
    system_prompt = config.get("system_prompt", "You are a helpful assistant.")

    if not llm_api_key or not website:
        raise ValueError("Please set sambanova_api_key and website in config.yaml")

    embed_model = HuggingFaceEmbeddings(model_name=config.get("embedding_model"))
    corpus = build_rag_corpus(config, embed_model, website)
    client = SambaNova(
        api_key=llm_api_key,
        base_url="https://api.sambanova.ai/v1",
        timeout=30,
    )

    RESOURCE_STATE.update(
        config=config,
        website=website,
        system_prompt=system_prompt,
        embed_model=embed_model,
        corpus=corpus,
        client=client,
    )
    return RESOURCE_STATE


def answer_question(question: str):
    resources = init_resources()
    selected = retrieve_relevant_chunks(
        resources["corpus"],
        question,
        resources["embed_model"],
        top_k=4,
    )
    prompt = build_prompt(resources["system_prompt"], question, selected)
    raw_answer = ask_model(prompt, resources["client"])
    response = format_answer(raw_answer, selected)
    citations = "\n\n".join(
        [f"Chunk {i+1}: {chunk.text[:300]}..." for i, chunk in enumerate(selected)]
    )
    return response, citations


def main():
    resources = init_resources()

    with gr.Blocks(title="RAG Chatbot") as demo:
        gr.Markdown("# 🤖 RAG-Powered Chatbot")
        gr.Markdown(f"**Website:** {resources['website']}  \n**Chunks:** {len(resources['corpus'])}")

        with gr.Row():
            with gr.Column(scale=3):
                question_input = gr.Textbox(label="Ask a question", placeholder="What services do you provide?", lines=2)
                submit_button = gr.Button("Ask")
                answer_output = gr.Textbox(label="Answer", lines=12, interactive=False)

            with gr.Column(scale=1):
                citations_output = gr.Textbox(label="Citations", lines=20, interactive=False)

        submit_button.click(
            answer_question,
            inputs=[question_input],
            outputs=[answer_output, citations_output],
        )

    demo.launch()


if __name__ == "__main__":
    main()