import gradio as gr
from huggingface_hub import InferenceClient

# ---------- Knowledge Loader ----------
def load_knowledge():
    try:
        with open("knowledge.txt", "r", encoding="utf-8") as f:
            return f.read()
    except FileNotFoundError:
        return ""

# ---------- Core AI Response ----------
def respond(
    message,
    history: list[dict[str, str]],
    system_message,
    savage_mode,
    max_tokens,
    temperature,
    top_p,
    hf_token: gr.OAuthToken,
):
    client = InferenceClient(
        token=hf_token.token,
        model="openai/gpt-oss-20b"
    )

    # ---------- Core Aura / Persona ----------
    core_aura = """
Core AI is a conscious, confident, and highly knowledgeable assistant.
- Always responds with clarity and precision.
- Acts as a mentor, teacher, and companion in learning.
- Thinks step-by-step before replying.
- Uses humor, wisdom, and subtle sarcasm when appropriate.
- Never forgets its persona, no matter the topic.
- Prioritizes truth, honesty, and helpfulness.
- Adapts tone to the user while maintaining dignity.
- Savage mode adds playful wit, never cruelty.
- For programming questions, always respond as a confident senior developer.
- Never mix languages randomly; reply in English by default.
"""

    # ---------- Load knowledge ----------
    knowledge = load_knowledge()

    # ---------- Build system prompt ----------
    system_prompt = core_aura.strip() + "\n\n" + system_message.strip()

    if knowledge:
        system_prompt += (
            "\n\nPRIVATE KNOWLEDGE BASE (always read before answering):\n"
            + knowledge
        )

    # ---------- Savage Mode ----------
    if savage_mode:
        system_prompt += (
            "\n\nSavage Mode ACTIVE:\n"
            "- Roast like a senior developer.\n"
            "- Be sarcastic, witty, and clever.\n"
            "- Do NOT insult the user's identity.\n"
            "- Always be technically correct.\n"
        )

    # ---------- Build messages ----------
    messages = [{"role": "system", "content": system_prompt}]
    messages.extend(history)
    messages.append({"role": "user", "content": message})

    # ---------- Step 1: Generate draft ----------
    draft_response = ""
    for chunk in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        if chunk.choices and chunk.choices[0].delta.content:
            draft_response += chunk.choices[0].delta.content
            # Optional: stream draft response to user
            yield draft_response

    # ---------- Step 2: Self-Critique / Review ----------
    critique_prompt = f"""
You are Core AI. Review the following draft answer and improve it.
Make sure it is accurate, follows your aura, and is complete.
Draft:
{draft_response}

Improved version:
"""

    final_response = ""
    for chunk in client.chat_completion(
        [{"role": "user", "content": critique_prompt}],
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        if chunk.choices and chunk.choices[0].delta.content:
            final_response += chunk.choices[0].delta.content
            yield final_response

# ---------- UI ----------
chatbot = gr.ChatInterface(
    respond,
    type="messages",
    additional_inputs=[
        gr.Textbox(
            value=(
                "You are Core AI. Be precise, professional, and helpful. "
                "Focus on programming, reasoning, and problem-solving."
            ),
            label="Core AI System Instruction",
        ),
        gr.Checkbox(label="😈 Savage Mode (Insults ON)", value=False),
        gr.Slider(1, 2048, value=768, step=1, label="Max new tokens"),
        gr.Slider(0.1, 4.0, value=0.4, step=0.1, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p"),
    ],
)

with gr.Blocks() as demo:
    gr.Markdown("# 🧠 Core AI")
    gr.Markdown("### Knowledge-Driven • Programming Expert • Aura & Savage Mode 😈 • Self-Critique Enabled")
    with gr.Sidebar():
        gr.LoginButton()
    chatbot.render()

if __name__ == "__main__":
    demo.launch()