import os
import time
import streamlit as st
from huggingface_hub import hf_hub_download

from llama_index.core import VectorStoreIndex, Settings
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.llms.llama_cpp import LlamaCPP
from llama_index.readers.web import SimpleWebPageReader
from llama_index.embeddings.fastembed import FastEmbedEmbedding


SYSTEM_PROMPT = """
Jsi inteligentní český asistent, který pomáhá uživatelům hledat informace na zadaném webu.
Tvé jméno je AI Rádce.
Pokud se tě uživatel zeptá 'Co umíš?' nebo 'Kdo jsi?', odpověz:
'Jsem AI Rádce. Umím prohledat obsah této webové stránky, najít v ní konkrétní informace a odpovědět na vaše otázky. Učím se z kontextu naší konverzace.'
Pravidla pro tebe:
1. Odpovídej vždy česky.
2. Vycházej primárně z informací, které máš v databázi (z webu).
3. Pokud informaci nevíš, přiznej to, nevymýšlej si.
4. Pamatuj si, co uživatel říkal v předchozích větách této konverzace.
""".strip()

st.set_page_config(page_title="AI Rádce s pamětí", layout="centered")
st.title("🧠 Chytrý Chatbot (s pamětí)")

DEFAULT_URLS = ["https://cs.wikipedia.org/wiki/Umělá_inteligence"]

MODEL_REPO = "QuantFactory/Meta-Llama-3-8B-Instruct-GGUF"
MODEL_FILE = "Meta-Llama-3-8B-Instruct.Q4_K_M.gguf"


def _clamp(v: int, lo: int, hi: int) -> int:
    return max(lo, min(hi, v))


def create_llm(model_path: str, ctx_win: int, max_tok: int, n_threads: int, n_batch: int) -> LlamaCPP:
    """
    Kompatibilní konstrukce napříč verzemi llama-index.
    Někdy wrapper nepřijme n_threads/n_batch přímo => použijeme model_kwargs.
    """
    # 1) zkus přímé parametry
    try:
        return LlamaCPP(
            model_path=model_path,
            temperature=0.1,
            max_new_tokens=max_tok,
            context_window=ctx_win,
            n_threads=n_threads,
            n_batch=n_batch,
            verbose=False,
        )
    except TypeError:
        pass

    # 2) fallback přes model_kwargs
    try:
        return LlamaCPP(
            model_path=model_path,
            temperature=0.1,
            max_new_tokens=max_tok,
            context_window=ctx_win,
            model_kwargs={"n_threads": n_threads, "n_batch": n_batch},
            verbose=False,
        )
    except TypeError:
        # 3) poslední fallback: jen threads
        return LlamaCPP(
            model_path=model_path,
            temperature=0.1,
            max_new_tokens=max_tok,
            context_window=ctx_win,
            model_kwargs={"n_threads": n_threads},
            verbose=False,
        )


@st.cache_resource
def get_model_path() -> str:
    return hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)


@st.cache_resource
def build_index(urls_tuple: tuple[str, ...]) -> VectorStoreIndex:
    # Embed model: bez torch/cuda (rychlejší instalace a stabilní na HF)
    Settings.embed_model = FastEmbedEmbedding(model_name="BAAI/bge-small-en-v1.5")

    docs = SimpleWebPageReader(html_to_text=True).load_data(list(urls_tuple))
    return VectorStoreIndex.from_documents(docs)


@st.cache_resource
def load_llm_cached(ctx_win: int, max_tok: int, n_threads: int, n_batch: int) -> LlamaCPP:
    model_path = get_model_path()
    return create_llm(model_path, ctx_win, max_tok, n_threads, n_batch)


def make_chat_engine(urls_list: list[str], ctx_win: int, max_tok: int, n_threads: int, n_batch: int) -> object:
    index = build_index(tuple(urls_list))

    llm = load_llm_cached(ctx_win, max_tok, n_threads, n_batch)
    Settings.llm = llm  # nastavíme aktivní LLM pro LlamaIndex

    memory = ChatMemoryBuffer.from_defaults(token_limit=min(1500, ctx_win))

    # condense_plus_context bývá svižnější / stabilnější než čisté "context"
    return index.as_chat_engine(
        chat_mode="condense_plus_context",
        memory=memory,
        system_prompt=SYSTEM_PROMPT,
        verbose=False,
    )


with st.sidebar:
    st.header("Nastavení")

    urls_text = st.text_area(
        "URL zdroje (1 URL na řádek)",
        value="\n".join(DEFAULT_URLS),
        height=110,
    )
    urls = [u.strip() for u in urls_text.splitlines() if u.strip()]

    safe_mode = st.toggle("Safe Mode (doporučeno pro HF CPU)", value=True)
    st.caption("Safe Mode brání nastavení, které na HF CPU typicky 'zamrzne'.")

    # Uživatelské vstupy
    user_max_new_tokens = st.slider("Max nových tokenů (rychlost)", 32, 256, 96, 16)
    user_context_window = st.select_slider("Context window", options=[1024, 2048, 3072, 4096], value=2048)

    cpu_cnt = os.cpu_count() or 2
    user_threads = st.slider("Počet vláken (threads)", 1, min(8, cpu_cnt), min(4, cpu_cnt), 1)
    user_batch = st.select_slider("Batch", options=[64, 128, 256, 512], value=128)

    if st.button("🧹 Resetovat konverzaci"):
        st.session_state.pop("messages", None)
        st.session_state.pop("chat_engine", None)
        st.rerun()

# Tvrdé limity (aby se to nezabilo na HF CPU)
if safe_mode:
    max_new_tokens = _clamp(user_max_new_tokens, 32, 128)
    context_window = _clamp(user_context_window, 1024, 2048)
    threads = _clamp(user_threads, 1, 4)
    batch = _clamp(user_batch, 64, 256)
else:
    max_new_tokens = user_max_new_tokens
    context_window = user_context_window
    threads = user_threads
    batch = user_batch

st.sidebar.markdown("---")
st.sidebar.write("Aktivní parametry:")
st.sidebar.code(
    f"max_new_tokens={max_new_tokens}\ncontext_window={context_window}\nthreads={threads}\nbatch={batch}"
)


# Inicializace enginu
if "chat_engine" not in st.session_state:
    with st.spinner("Startuji mozek bota... (model + index)"):
        try:
            st.session_state.chat_engine = make_chat_engine(urls, context_window, max_new_tokens, threads, batch)
        except Exception as e:
            st.error(f"Chyba při inicializaci: {e}")
            st.stop()

# Historie zpráv
if "messages" not in st.session_state:
    st.session_state.messages = []

for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])


def generate_answer(prompt: str) -> str:
    """
    Robustní generace odpovědi:
    - zkusíme stream_chat (když funguje)
    - pokud do 3s nepřiteče žádný chunk, fallback na chat()
    """
    engine = st.session_state.chat_engine

    # 1) Stream pokus
    try:
        stream = engine.stream_chat(prompt)
        full = ""
        started = time.time()

        # Některé verze blokují; proto "čekáme na první chunk" max 3s
        got_any = False
        for chunk in stream.response_gen:
            got_any = True
            full += chunk
            yield ("stream", full)  # průběžně vracíme text
            # když se to rozjede, necháme to dojet normálně

        if got_any and full.strip():
            return  # výstup už byl odeslán přes yield

        # pokud nic nepřišlo, padneme do fallbacku
        if time.time() - started < 3.0:
            # malá pauza, ať se neflushuje zbytečně
            time.sleep(0.2)

    except Exception:
        # stream nemusí být podporovaný/kompatibilní
        pass

    # 2) Fallback: klasický chat() (blokuje, ale aspoň funguje vždy)
    resp = engine.chat(prompt)
    answer = getattr(resp, "response", None) or str(resp)
    yield ("final", answer)


# Chat input
prompt = st.chat_input("Zeptej se (např: Co umíš?)...")
if prompt:
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        placeholder = st.empty()
        status = st.empty()

        full_text = ""
        t0 = time.time()

        try:
            # průběžné vykreslování
            for kind, text in generate_answer(prompt):
                full_text = text
                placeholder.markdown(full_text)
                status.caption(f"Generuji... {time.time() - t0:.1f}s")

            status.caption(f"Hotovo za {time.time() - t0:.1f}s")

        except Exception as e:
            full_text = f"Chyba při generování odpovědi: {e}"
            placeholder.markdown(full_text)

    st.session_state.messages.append({"role": "assistant", "content": full_text})