app/ui_gradio.py

import asyncio
import logging
import os
import random
import string
import traceback
import uuid
from typing import Dict, List, Tuple

from app.services import session_store as _ss

import gradio as gr
import httpx

from app.core.config import get_settings
from app.services.embeddings import EmbeddingService
from app.services.llm import LLMService
from app.services.rag_pipeline import RAGPipeline
from app.services.vector_store import FaissVectorStore
from app.services.reranker import RerankerService

logger = logging.getLogger(__name__)
settings = get_settings()
API_URL = os.getenv("RAG_API_URL", "").strip()

# Initialize services
embedding_service = EmbeddingService(settings.embedding_model)
vector_store = FaissVectorStore(
    embedding_service=embedding_service,
    docs_dir=settings.docs_dir,
    index_dir=settings.index_dir,
    chunk_size_tokens=settings.chunk_size_tokens,
    chunk_overlap_tokens=settings.chunk_overlap_tokens,
)
llm_service = LLMService(
    provider=settings.llm_provider,
    groq_api_key=settings.groq_api_key,
    groq_model=settings.groq_model,
    groq_rewrite_model=settings.groq_rewrite_model,
    hf_api_key=settings.hf_api_key,
    hf_model=settings.hf_model,
    timeout_s=settings.request_timeout_s,
)
reranker_service = RerankerService(settings.reranker_model)
pipeline = RAGPipeline(
    vector_store=vector_store, 
    llm_service=llm_service, 
    reranker=reranker_service,
    top_k=settings.top_k,
    max_context_chunks=settings.max_context_chunks
)

# Ensure index is ready
vector_store.build_or_load()


# ── Keyboard layout for realistic typos ──────────────────────────────
NEARBY_KEYS = {
    'a': 'sqwz', 'b': 'vghn', 'c': 'xdfv', 'd': 'sfecx', 'e': 'wrsdf',
    'f': 'dgrtcv', 'g': 'fhtybn', 'h': 'gjyunb', 'i': 'uojkl', 'j': 'hkunmi',
    'k': 'jlomi', 'l': 'kop', 'm': 'njk', 'n': 'bhjm', 'o': 'ipkl',
    'p': 'ol', 'q': 'wa', 'r': 'edft', 's': 'awedxz', 't': 'rfgy',
    'u': 'yihj', 'v': 'cfgb', 'w': 'qase', 'x': 'zsdc', 'y': 'tghu',
    'z': 'xas',
}

TYPO_CHANCE = 0.07  # 7% chance per word


def _nearby_char(ch: str) -> str:
    """Return a plausible neighbouring key for the given character."""
    lower = ch.lower()
    if lower in NEARBY_KEYS:
        replacement = random.choice(NEARBY_KEYS[lower])
        return replacement.upper() if ch.isupper() else replacement
    return random.choice(string.ascii_lowercase)


def _to_history(messages: List[Dict[str, str]]) -> list:
    """Returns the history as a list of dicts (already in this format for Gradio 5)."""
    return messages


async def _chat_via_api(message: str, history: List[Dict[str, str]]) -> str:
    payload = {"message": message, "history": history}
    headers = {"Content-Type": "application/json"}
    if settings.api_key:
        headers["x-api-key"] = settings.api_key
    
    async with httpx.AsyncClient(timeout=30.0) as client:
        resp = await client.post(API_URL, json=payload, headers=headers)
        if resp.status_code == 429:
            return "⚠️ I'm receiving too many requests right now. Please wait a moment before sending another message."
        resp.raise_for_status()
        return resp.json()["reply"]


async def _get_reply(message: str, chat_history: List[Dict[str, str]]) -> str:
    """Fetch the full reply from the LLM (API or direct pipeline)."""
    if API_URL:
        return await _chat_via_api(message, chat_history)
    result = await pipeline.chat(message=message, history=chat_history)
    return result["reply"]


async def chat_fn(message: str, chat_history: List[Dict[str, str]], session_id: str = ""):
    """Streaming generator: yields word-by-word with human-like timing & typos."""
    if not message or not message.strip():
        yield gr.update(interactive=True), chat_history
        return

    original_history = chat_history.copy()
    chat_history = chat_history + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": ""}
    ]
    yield gr.update(value="", interactive=False), chat_history

    try:
        reply = await _get_reply(message, original_history)
    except httpx.HTTPStatusError as e:
        logger.error(f"HTTP Error: {e.response.status_code} - {e.response.text}")
        reply = "⚠️ Rate limit reached. Please slow down a bit!" if e.response.status_code == 429 \
            else "⚠️ I encountered an error. Please try again in a few seconds."
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}\n{traceback.format_exc()}")
        reply = f"⚠️ Oops! Something went wrong."

    # ── Humanistic letter-by-letter typing — HTML-safe ────────────────────
    displayed = ""       # what is shown in the chat bubble
    buffer = ""          # invisible accumulator for mid-HTML-tag characters
    in_tag = False       # True while we are inside a < … > sequence

    for char in reply:
        if char == '<':
            in_tag = True
            buffer += char
            continue

        if in_tag:
            buffer += char
            if char == '>':
                # Tag is now complete — flush it to the display all at once
                in_tag = False
                displayed += buffer
                buffer = ""
                chat_history[-1] = {"role": "assistant", "content": displayed}
                yield gr.update(value="", interactive=False), chat_history
                # Small pause after a <br> to mimic line-break pacing
                if displayed.endswith('<br>') or displayed.endswith('<br/>'):
                    await asyncio.sleep(random.uniform(0.1, 0.2))
            continue

        # ── Normal character (not inside a tag) ───────────────────────────
        # Decide typo chance only for plain alphabetic words
        do_typo = (
            char.isalpha()
            and len(displayed) > 8           # skip the very beginning
            and random.random() < TYPO_CHANCE
        )

        if do_typo:
            wrong_char = _nearby_char(char)
            displayed += wrong_char
            chat_history[-1] = {"role": "assistant", "content": displayed}
            yield gr.update(value="", interactive=False), chat_history
            await asyncio.sleep(random.uniform(0.15, 0.3))

            # Backspace the wrong character
            displayed = displayed[:-1]
            chat_history[-1] = {"role": "assistant", "content": displayed}
            yield gr.update(value="", interactive=False), chat_history
            await asyncio.sleep(random.uniform(0.1, 0.2))

        # Type the correct character
        displayed += char
        chat_history[-1] = {"role": "assistant", "content": displayed}
        yield gr.update(value="", interactive=False), chat_history

        # Pacing: punctuation pauses, space micro-pause, normal char speed
        if char in '.!?':
            await asyncio.sleep(random.uniform(0.35, 0.7))
        elif char in ',:;':
            await asyncio.sleep(random.uniform(0.15, 0.3))
        elif char == ' ':
            await asyncio.sleep(random.uniform(0.06, 0.13))
        else:
            await asyncio.sleep(random.uniform(0.02, 0.06))

    # ── Fire-and-forget: persist to session store (never blocks the UI) ───
    # We only save successful responses, not error messages
    if session_id and not displayed.startswith("⚠️"):
        asyncio.create_task(_ss.save_message(session_id, message, displayed))

    # Re-enable the input box at the end
    yield gr.update(interactive=True), chat_history


custom_css = """
/* Aggressively force height: auto and min-height: 40px on the chatbot and flex containers */
.gradio-container, .flex, .block, #chatbot-window {
    height: auto !important;
    min-height: 40px !important;
}

/* Ensure the chatbot doesn't overflow the viewport if it gets too full */
#chatbot-window {
    max-height: 70vh !important;
    border: none !important;
}

/* Aggressively hide the footer, including HF injected footers if inside the container */
footer, .footer, footer * {
    display: none !important;
    visibility: hidden !important;
    height: 0 !important;
    margin: 0 !important;
    padding: 0 !important;
}
"""

# JS: runs immediately on page load — generates/loads session_id from localStorage
# and restores previous chat history. Completely non-blocking.
_SESSION_JS = """
async () => {
    // ── Session ID: generate once, persist forever in localStorage ──
    let sid = localStorage.getItem('martech_session_id');
    if (!sid) {
        sid = 'sess-' + Date.now() + '-' + Math.random().toString(36).slice(2, 9);
        localStorage.setItem('martech_session_id', sid);
    }
    // Write the session_id into the hidden Gradio textbox
    const el = document.getElementById('session-id-box')?.querySelector('textarea');
    if (el) { el.value = sid; el.dispatchEvent(new Event('input', {bubbles:true})); }
    return sid;
}
"""

with gr.Blocks(title="Martechsol Assistant", theme=gr.themes.Soft(), css=custom_css) as demo:
    gr.Markdown("# Martechsol Assistant")
    gr.Markdown("Welcome! How can I help you today?")
    
    chatbot = gr.Chatbot(
        show_label=False,
        elem_id="chatbot-window",
        render_markdown=True,
    )
    with gr.Row():
        msg = gr.Textbox(
            placeholder="Type your question here...",
            show_label=False,
            scale=9
        )
        send = gr.Button("Send", variant="primary", scale=1)
    
    clear = gr.Button("Clear Chat History")

    # Hidden textbox holds the session_id — invisible to the user
    session_id = gr.Textbox(
        value="",
        visible=False,
        elem_id="session-id-box",
        label="session_id",
    )

    # On load: run JS to set session_id from localStorage (non-blocking)
    demo.load(fn=None, inputs=None, outputs=session_id, js=_SESSION_JS)

    # Wire up the events — streaming generator for live typing effect (UNCHANGED)
    send.click(chat_fn, inputs=[msg, chatbot, session_id], outputs=[msg, chatbot])
    msg.submit(chat_fn, inputs=[msg, chatbot, session_id], outputs=[msg, chatbot])
    clear.click(lambda: [], None, chatbot)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)