import base64, os, uuid, gradio as gr
from pathlib import Path
from typing import Dict, List, Tuple
from langfuse import Langfuse, observe, get_client
from langfuse.openai import openai

langfuse = Langfuse(secret_key=os.getenv('LANGFUSE_SECRET_KEY'), public_key=os.getenv('LANGFUSE_PUBLIC_KEY'), host=os.getenv('LANGFUSE_HOST'))
langfuse = get_client()

BASE_URL = os.getenv('CHAT_BASE_URL')
MODEL_NAME = os.getenv('MODEL_NAME')
API_KEY = os.getenv('API_KEY')
SYSTEM_PROMPT = (os.getenv('SYSTEM_PROMPT'))

LOGO_PATH = Path(__file__).parent / "logo.png"

def _load_logo_data_uri() -> str:
    try:
        encoded = base64.b64encode(LOGO_PATH.read_bytes()).decode("ascii")
        return f"data:image/png;base64,{encoded}"
    except FileNotFoundError: return ""

LOGO_DATA_URI = _load_logo_data_uri()
HERO_LOGO_HTML = (
    f'<div class="hero-logo"><img src="{LOGO_DATA_URI}" alt="OpenJAI logo" /></div>'
    if LOGO_DATA_URI
    else ""
)

QUICK_PROMPTS: List[str] = ["Who is JTS?", "What is the benefit of small LLM?", "คุณเป็นใคร", "AI ทำงานยังไง",]

HERO_SECTION = f"""
<div class="hero-banner">
  <div class="hero-content">
    {HERO_LOGO_HTML}
    <div class="hero-copy">
      <span class="badge">Introducing OpenJAI-v1.0</span>
      <p>Partner with OpenJAI to deliver exceptional Thai language experiences, manage long-form context, and power reliable RAG-ready workflows. Purpose-built by the JTS research team.</p>
      <div class="hero-highlights">
        <span>⚡ Optimized for Thai and English</span>
        <span>🧠 Resilient in long-context RAG</span>
        <span>🧩 Compact & efficient</span>
      </div>
    </div>
  </div>
</div>
"""

INSIGHTS_HTML = """
<div class="surface-card">
  <h3>Why OpenJAI?</h3>
  <ul class="feature-grid">
    <li>
      <span class="feature-icon">🧡</span>
      <div>
        <h4>Better Thai and English text comprehension</h4>
      </div>
    </li>
    <li>
      <span class="feature-icon">🪄</span>
      <div>
        <h4>RAG-focused</h4>
        <p>Stronger understanding of long-context inputs.</p>
      </div>
    </li>
    <li>
      <span class="feature-icon">🛡️</span>
      <div>
        <h4>Robust & lightweight</h4>
        <p>Compact 14B model designed for efficiency with a strong performance.</p>
      </div>
    </li>
  </ul>
</div>
<div class="surface-card secondary">
  <h3>Learn more at</h3>
  <ul class="moments-list">
    <li>
        <a href="https://arxiv.org/abs/2510.06847" target="_blank" rel="noopener noreferrer" style="color: white;">
            <strong>Technical Report</strong>
        </a>
    </li>
  </ul>
</div>
"""

FOOTER_HTML = """
<div class="footer-note">
  Crafted for Thai innovators by <strong>JTS</strong>. Share your ideas with us at <a href="mailto:jts.ai.team@gmail.com">jts.ai.team@gmail.com</a>.
</div>
"""

CUSTOM_CSS = """:root,
[data-theme="light"],
[data-theme="dark"] {
  --bg-primary: #1f130b;
  --bg-gradient-1: rgba(255, 173, 96, 0.22);
  --bg-gradient-2: rgba(255, 113, 91, 0.2);
  --text-primary: #fff3e6;
  --text-secondary: rgba(255, 244, 235, 0.88);
  --text-tertiary: rgba(255, 234, 214, 0.82);
  
  --hero-bg-1: rgba(255, 159, 67, 0.28);
  --hero-bg-2: rgba(255, 98, 0, 0.32);
  --hero-border: rgba(255, 240, 225, 0.24);
  --hero-shadow: rgba(110, 48, 0, 0.45);
  
  --card-bg: rgba(54, 24, 6, 0.78);
  --card-bg-secondary: rgba(64, 26, 8, 0.8);
  --card-border: rgba(255, 202, 149, 0.22);
  --card-shadow: rgba(48, 20, 4, 0.5);
  
  --chat-bg: rgba(43, 21, 7, 0.86);
  --chat-border: rgba(255, 204, 152, 0.2);
  --message-bot-bg: rgba(255, 199, 150, 0.08);
  --message-user-bg: rgba(255, 142, 78, 0.22);
  
  --input-bg: rgba(36, 18, 9, 0.72);
  --input-border: rgba(255, 198, 150, 0.18);
  --input-focus: rgba(255, 183, 92, 0.65);
  
  --badge-bg: rgba(255, 255, 255, 0.16);
  --highlight-bg: rgba(92, 45, 7, 0.6);
  --highlight-border: rgba(255, 205, 169, 0.32);
}

:root,
[data-theme="light"],
[data-theme="dark"] {
  --block-background-fill: transparent;
  --body-background-fill: #1f130b;
  --background-fill-primary: #1f130b;
  --body-text-color: #fff3e6;

body {
  margin: 0;
  background: radial-gradient(circle at 20% 20%, var(--bg-gradient-1), transparent 45%),
              radial-gradient(circle at 85% 10%, var(--bg-gradient-2), transparent 40%),
              var(--bg-primary);
  color: var(--text-primary);
  font-family: 'Inter', 'Segoe UI', system-ui, -apple-system, sans-serif;
}

.gradio-container {
  max-width: 1200px;
  margin: 0 auto;
  padding: 40px 28px 80px !important;
  background: transparent;
}

.hero-banner {
  position: relative;
  padding: 52px 48px 56px;
  border-radius: 28px;
  background: linear-gradient(135deg, var(--hero-bg-1), var(--hero-bg-2));
  backdrop-filter: blur(18px);
  border: 1px solid var(--hero-border);
  overflow: hidden;
  box-shadow: 0 36px 120px var(--hero-shadow);
}

.hero-banner::after {
  content: "";
  position: absolute;
  inset: 0;
  background: radial-gradient(circle at 18% 22%, rgba(255, 255, 255, 0.18), transparent 55%);
  pointer-events: none;
}

.hero-content {
  position: relative;
  display: flex;
  align-items: center;
  gap: 32px;
}

.hero-logo {
  flex: 0 0 auto;
  display: flex;
  align-items: center;
  justify-content: center;
  padding: 18px;
  border-radius: 24px;
  background: var(--badge-bg);
  border: 1px solid var(--hero-border);
  box-shadow: 0 18px 50px var(--card-shadow);
}

.hero-logo img {
  width: 120px;
  height: auto;
  display: block;
}

.hero-copy {
  flex: 1 1 auto;
  display: flex;
  flex-direction: column;
  gap: 16px;
  min-width: 0;
}

.hero-banner h1 {
  font-size: 48px;
  line-height: 1.05;
  margin: 18px 0 18px;
  color: var(--text-primary);
}

.hero-banner p {
  max-width: 640px;
  font-size: 19px;
  color: var(--text-secondary);
  margin-bottom: 26px;
}

.badge {
  display: inline-flex;
  align-items: center;
  gap: 8px;
  background: var(--badge-bg);
  border-radius: 999px;
  padding: 6px 16px;
  font-size: 14px;
  letter-spacing: 0.4px;
  text-transform: uppercase;
  color: var(--text-primary);
  font-weight: 600;
}

.hero-highlights {
  display: flex;
  flex-wrap: wrap;
  gap: 12px;
}

.hero-highlights span {
  display: inline-flex;
  align-items: center;
  gap: 10px;
  padding: 10px 18px;
  border-radius: 999px;
  background: var(--highlight-bg);
  border: 1px solid var(--highlight-border);
  font-size: 15px;
  color: var(--text-primary);
  font-weight: 500;
}

.content-row {
  margin-top: 42px;
  gap: 24px;
}

.chat-column {
  gap: 16px;
}

.chat-panel {
  border-radius: 24px !important;
  background: var(--chat-bg) !important;
  border: 1px solid var(--chat-border) !important;
  box-shadow: 0 28px 80px var(--card-shadow);
}

.chat-panel .message.bot {
  background: var(--message-bot-bg) !important;
  border-radius: 18px;
  color: var(--text-primary) !important;
}

.chat-panel .message.user {
  background: var(--message-user-bg) !important;
  border-radius: 18px;
  color: var(--text-primary) !important;
}

.compose-panel {
  padding: 10px;
  border-radius: 20px;
  # background: var(--chat-bg) !important;
  border: 1px solid var(--card-border);
  # box-shadow: inset 0 0 0 1px var(--input-border);
  gap: 0px;
}

.compose-panel textarea {
  background: var(--chat-bg) !important;
  border: 1px solid var(--input-border);
  border-radius: 16px;
  font-size: 16px;
  min-height: 94px !important;
}

.compose-panel textarea:focus {
  border-color: black;
  box-shadow: 0 0 0 3px rgba(255, 162, 53, 0.22);
}

.actions-row {
  display: flex;
  justify-content: space-between;
  gap: 14px;
}

.primary-btn {
  background: linear-gradient(135deg, #ff9933, #ff6a2b) !important;
  border: none !important;
  color: #fff !important;
  font-weight: 600 !important;
  padding: 12px 26px !important;
  border-radius: 14px !important;
}

.primary-btn:hover {
  filter: brightness(1.08);
}

.ghost-btn {
  background: linear-gradient(135deg, #ff9933, #ff6a2b) !important;
  border: none !important;
  color: var(--text-primary) !important;
  font-weight: 500 !important;
  padding: 12px 26px !important;
  border-radius: 14px !important;
}

.quick-prompts {
  margin-top: 10px;
  padding: 24px;
  border: 1px solid var(--input-border);
  border-radius: 16px;
  # background-color: transparent !important;
}

.quick-prompts .label {
  font-size: 17px;
  color: #CC7000;
  # color: var(--text-tertiary);
  # background-color: transparent !important;
  margin-bottom: 14px;
  padding: 0;
  font-weight: 600;
}

.quick-grid {
  display: grid;
  grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
  gap: 12px;
}

.quick-chip {
  justify-content: flex-start !important;
  border-radius: 14px !important;
  background: var(--card-bg) !important;
  border: 1px solid var(--card-border) !important;
  color: var(--text-primary) !important;
  font-weight: 500 !important;
  min-height: 54px !important;
}

.quick-chip:hover {
  border-color: var(--input-focus) !important;
  transform: translateY(-1px);
}

.info-column {
  display: flex;
  flex-direction: column;
  gap: 18px;
}

.surface-card {
  padding: 28px;
  border-radius: 24px;
  background: var(--card-bg);
  border: 1px solid var(--card-border);
  box-shadow: 0 24px 70px var(--card-shadow);
  margin-bottom: 24px;
}

.surface-card.secondary {
  background: var(--card-bg-secondary);
}

.surface-card h3 {
  margin-top: 0;
  font-size: 22px;
  color: var(--text-primary);
}

.feature-grid {
  list-style: none;
  padding: 0;
  margin: 18px 0 0;
  display: flex;
  flex-direction: column;
  gap: 18px;
}

.feature-grid li {
  display: flex;
  gap: 18px;
  align-items: flex-start;
}

.feature-icon {
  font-size: 26px;
  line-height: 1;
}

.feature-grid h4 {
  margin: 0 0 4px;
  font-size: 18px;
  color: var(--text-primary);
}

.feature-grid p {
  margin: 0;
  color: var(--text-secondary);
  line-height: 1.55;
}

.moments-list {
  list-style: none;
  padding: 0;
  margin: 16px 0 0;
  display: grid;
  gap: 12px;
}

.moments-list li {
  color: var(--text-secondary);
  line-height: 1.55;
}

.footer-note {
  margin-top: 48px;
  text-align: center;
  color: var(--text-tertiary);
  font-size: 15px;
}

.footer-note a {
  color: #ff9933;
  font-weight: 500;
  text-decoration: none;
}

@media (max-width: 992px) {
  .hero-banner {
    padding: 36px 28px;
  }
  .hero-banner h1 {
    font-size: 38px;
  }
  .hero-content {
    flex-direction: column;
    align-items: flex-start;
    gap: 24px;
  }
  .hero-logo {
    padding: 16px;
  }
  .content-row {
    flex-direction: column;
  }
}

@media (max-width: 640px) {
  .hero-banner {
    padding: 28px 24px;
  }
  .hero-content {
    gap: 18px;
  }
  .hero-highlights {
    flex-direction: column;
  }
}"""

try: client = openai.AsyncOpenAI(base_url=BASE_URL, api_key=API_KEY)
except Exception as exc:  
    print(f"Error initializing OpenAI client: {exc}")
    client = None

@observe()
async def get_api_stream(messages: List[Dict[str, str]], user_id: str):
    """
    This inner function is observed by Langfuse.
    It performs the API call and returns the raw stream object.
    """
    if langfuse: langfuse.update_current_trace(user_id=user_id)
    
    stream = await client.chat.completions.create(
        model=MODEL_NAME,
        messages=messages,
        extra_body={"chat_template_kwargs": {"enable_thinking": False}},
        stream=True,
        max_tokens=4096,
    )
    return stream


async def stream_chat(history: List[Tuple[str, str]], user_message: str, user_id: str):
    """
    This outer function is the generator that Gradio interacts with.
    It calls the observed inner function and yields updates to the UI.
    """
    history = history or []
    if not user_message or not user_message.strip():
        yield history, history
        return

    working_history = history + [(user_message, "")]
    yield working_history, working_history

    if not client:
        working_history[-1] = (user_message, "❌ Error: AI client not configured.")
        yield working_history, working_history
        return

    messages: List[Dict[str, str]] = [{"role": "system", "content": SYSTEM_PROMPT}]
    for previous_user, previous_bot in history:
        messages.append({"role": "user", "content": previous_user})
        messages.append({"role": "assistant", "content": previous_bot})
    messages.append({"role": "user", "content": user_message})

    assistant_text = ""
    try:
        stream = await get_api_stream(messages, user_id)
        async for chunk in stream:
            delta = chunk.choices[0].delta.content
            if delta:
                assistant_text += delta
                working_history[-1] = (user_message, assistant_text)
                yield working_history, working_history
                
    except Exception as error:
        working_history[-1] = (user_message, f"⚠️ Error: {error}")
        yield working_history, working_history

def reset_conversation(): return [], [], ""

with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="orange"),
    title="OpenJAI-v1.0 · Modern AI Companion",
    css=CUSTOM_CSS,
) as demo:
    gr.HTML(HERO_SECTION)

    with gr.Row(elem_classes="content-row"):
        with gr.Column(scale=7, elem_classes="chat-column"):
            chatbot = gr.Chatbot(
                label="Conversation history",
                bubble_full_width=False,
                height=520,
                show_copy_button=True,
                elem_classes="chat-panel",
            )

            with gr.Group(elem_classes="compose-panel"):
                prompt_input = gr.Textbox(
                    show_label=False,
                    placeholder="Type your message here...",
                    lines=4,
                    autofocus=True,
                )
                with gr.Row(elem_classes="actions-row"):
                    submit_button = gr.Button("Send", elem_classes="primary-btn")
                    clear_button = gr.Button(
                        "Reset Chat", elem_classes="ghost-btn", variant="secondary"
                    )

        with gr.Column(scale=5, min_width=320, elem_classes="info-column"):
            gr.HTML(INSIGHTS_HTML)

    with gr.Group(elem_classes="quick-prompts"):
        gr.HTML("<div class='label'>Suggest question:</div>")
        with gr.Row(elem_classes="quick-grid"):
            quick_buttons = []
            for prompt in QUICK_PROMPTS:
                button = gr.Button(prompt, elem_classes="quick-chip", variant="secondary")
                quick_buttons.append((button, prompt))

    gr.HTML(FOOTER_HTML)

    def create_session_uuid(): return str(uuid.uuid4())
    
    history_state = gr.State([])
    user_id_state = gr.State()

    demo.load(fn=create_session_uuid, inputs=None, outputs=user_id_state)

    submit_event = submit_button.click(
        fn=stream_chat,
        inputs=[history_state, prompt_input, user_id_state],
        outputs=[chatbot, history_state],
    )
    submit_event.then(fn=lambda: "", inputs=None, outputs=prompt_input, queue=False)

    submit_input_event = prompt_input.submit(
        fn=stream_chat,
        inputs=[history_state, prompt_input, user_id_state],
        outputs=[chatbot, history_state],
    )
    submit_input_event.then(fn=lambda: "", inputs=None, outputs=prompt_input, queue=False)

    clear_button.click(
        fn=reset_conversation,
        inputs=None,
        outputs=[chatbot, history_state, prompt_input],
        queue=False,
    )

    for button, prompt in quick_buttons:
        button.click(fn=lambda p=prompt: p, inputs=None, outputs=prompt_input, queue=False)

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