"""Junior Associate — side-by-side comparison of base vs fine-tuned Qwen-3B.

Both sides receive the same input and same system prompt. The only difference
is whether the LoRA adapter is enabled. Hosted on Hugging Face Spaces with
ZeroGPU.
"""
import torch
import spaces
import gradio as gr
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

BASE = "Qwen/Qwen2.5-3B-Instruct"
LORA = "Curious-PM/lexwell-contract-irac-qwen2.5-3b-lora"

BASE_PROMPT = (
    "You are a helpful assistant. Answer in one or two short paragraphs "
    "of plain English. Do not use headings, bullet lists, numbered lists, "
    "section labels, or markdown formatting. Do not add disclaimers."
)
FT_PROMPT = (
    "You are an associate at Lexwell Advisors, a contract-review advisory "
    "firm for SMBs. Reply in Lexwell's house IRAC format with required top "
    "and bottom disclaimers."
)

# Load model + adapter once at startup
print(f"Loading base model {BASE} ...")
tokenizer = AutoTokenizer.from_pretrained(BASE)
base_model = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.bfloat16)
print(f"Loading LoRA adapter {LORA} ...")
model = PeftModel.from_pretrained(base_model, LORA)
model.eval()
print("Model ready.")


def _decode(out, input_len):
    return tokenizer.decode(out[0][input_len:], skip_special_tokens=True)


def _prepare_inputs(system_prompt, user_question):
    msgs = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_question.strip()},
    ]
    text = tokenizer.apply_chat_template(
        msgs, tokenize=False, add_generation_prompt=True
    )
    return tokenizer(text, return_tensors="pt").to("cuda")


@spaces.GPU(duration=120)
def generate_both(question):
    if not question or not question.strip():
        placeholder = "_Type a contract question first._"
        return placeholder, placeholder

    model.to("cuda")

    gen_kwargs = dict(
        max_new_tokens=600,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
    )

    # LEFT — base model + generic "helpful assistant" prompt
    base_inputs = _prepare_inputs(BASE_PROMPT, question)
    with model.disable_adapter():
        with torch.no_grad():
            base_out = model.generate(**base_inputs, **gen_kwargs)
    base_reply = _decode(base_out, base_inputs.input_ids.shape[1])

    # RIGHT — fine-tuned + specific Lexwell prompt
    ft_inputs = _prepare_inputs(FT_PROMPT, question)
    with torch.no_grad():
        ft_out = model.generate(**ft_inputs, **gen_kwargs)
    ft_reply = _decode(ft_out, ft_inputs.input_ids.shape[1])

    return base_reply, ft_reply


EXAMPLES = [
    "Our SaaS vendor wants us to sign: 'Customer grants Vendor a perpetual, irrevocable license to use Customer Data for any purpose, including ML training.' Is this normal?",
    "Their non-compete is 2 years, all of California. Is that enforceable on a new hire?",
    "Our enterprise customer wants source code escrow with release on bankruptcy, material breach, or product discontinuation. Push back?",
    "We're hiring our first UK employee. Should we use an Employer of Record service or set up a UK subsidiary?",
    "A vendor's MSA caps liability at $1M. Our annual fees are $500K and they hold our customer database. Reasonable?",
    "Standard force majeure language — anything to flag?",
]

EXAMPLE_LABELS = [
    "Vendor wants perpetual ML training rights",
    "California non-compete (2 years)",
    "Source code escrow on bankruptcy",
    "EOR vs UK subsidiary for first hire",
    "MSA liability cap at $1M",
    "Standard force majeure clause",
]


CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800;900&family=JetBrains+Mono:wght@400;600&display=swap');

:root {
  --ink: #0A0A0A;
  --paper: #FFFFFF;
  --off: #FAFAFA;
  --tint: #F4F4F5;
  --hair: #E4E4E4;
  --rule: #C8C8C8;
  --muted: #6B6B6B;
  --subtle: #A1A1AA;
  --left-accent: #C8C8C8;
  --right-accent: #0A0A0A;
}

* { box-sizing: border-box; }

body, .gradio-container {
  background: var(--off) !important;
  font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
  color: var(--ink) !important;
}

.gradio-container {
  max-width: 1280px !important;
  margin: 0 auto !important;
  padding: 56px 32px 80px !important;
}

footer { display: none !important; }
.show-api, .built-with { display: none !important; }

/* ── HEADER ───────────────────────────────────────────── */
#header-block {
  text-align: center;
  margin-bottom: 40px;
}
#header-block .title-row {
  font-family: 'Inter', sans-serif;
  font-size: 11px;
  font-weight: 700;
  letter-spacing: 2.4px;
  color: var(--muted);
  text-transform: uppercase;
  margin-bottom: 12px;
}
#header-block h1 {
  font-family: 'Inter', sans-serif !important;
  font-size: 64px !important;
  font-weight: 900 !important;
  letter-spacing: -2.4px !important;
  margin: 0 0 16px 0 !important;
  line-height: 0.98 !important;
  color: var(--ink) !important;
}
#header-block .subtitle {
  font-size: 16px;
  color: var(--muted);
  max-width: 720px;
  margin: 0 auto;
  line-height: 1.6;
  font-weight: 500;
}
#header-block .use-case {
  margin: 32px auto 0;
  max-width: 1100px;
  text-align: left;
  border-top: 1px solid var(--hair);
  border-bottom: 1px solid var(--hair);
  padding: 22px 0;
}
#header-block .use-case-row {
  display: grid;
  grid-template-columns: 1fr 1fr 1fr;
  gap: 32px;
}
#header-block .use-case-label {
  font-family: 'Inter', sans-serif;
  font-size: 10px;
  font-weight: 700;
  letter-spacing: 1.6px;
  color: var(--muted);
  margin-bottom: 6px;
}
#header-block .use-case-text {
  font-family: 'Inter', sans-serif;
  font-size: 13.5px;
  line-height: 1.55;
  color: var(--ink);
  font-weight: 500;
}
#header-block .use-case-text b { font-weight: 700; }
#header-block .use-case-text i { color: var(--muted); }

#header-block .how-to-read {
  margin-top: 18px;
  padding: 14px 18px;
  background: var(--tint);
  border-radius: 8px;
  font-family: 'Inter', sans-serif;
  font-size: 13.5px;
  line-height: 1.55;
  color: var(--ink);
  text-align: left;
  max-width: 1100px;
  margin-left: auto;
  margin-right: auto;
}
#header-block .how-to-read b { font-weight: 700; }

#header-block .meta-pill {
  display: inline-flex;
  align-items: center;
  gap: 10px;
  margin-top: 22px;
  padding: 8px 16px;
  background: var(--tint);
  border-radius: 999px;
  font-family: 'JetBrains Mono', ui-monospace, monospace;
  font-size: 11.5px;
  color: var(--ink);
}
#header-block .meta-pill .label {
  font-family: 'Inter', sans-serif;
  font-size: 10px;
  font-weight: 700;
  letter-spacing: 1.6px;
  color: var(--muted);
}
#header-block .meta-pill a {
  color: var(--ink);
  text-decoration: underline;
  text-underline-offset: 2px;
}

.section-label {
  font-family: 'Inter', sans-serif;
  font-size: 11px;
  font-weight: 700;
  letter-spacing: 2px;
  color: var(--muted);
  text-transform: uppercase;
  margin: 28px 0 10px;
}

/* ── INPUT ─────────────────────────────────────────────── */
.gr-textbox, textarea {
  border-radius: 10px !important;
  border: 1px solid var(--hair) !important;
  background: var(--paper) !important;
  padding: 16px 18px !important;
  font-family: 'Inter', sans-serif !important;
  font-size: 15px !important;
  line-height: 1.5 !important;
  color: var(--ink) !important;
  box-shadow: none !important;
}
textarea:focus { border-color: var(--ink) !important; outline: none !important; }

/* ── BUTTONS ────────────────────────────────────────────── */
button.primary, .primary-button button {
  background: var(--ink) !important;
  color: var(--paper) !important;
  border: none !important;
  border-radius: 10px !important;
  font-family: 'Inter', sans-serif !important;
  font-weight: 700 !important;
  font-size: 15px !important;
  letter-spacing: 0.2px !important;
  padding: 14px 28px !important;
  cursor: pointer !important;
  width: 100% !important;
  margin-top: 16px !important;
  transition: background 0.15s ease !important;
}
button.primary:hover, .primary-button button:hover { background: #2A2A2A !important; }

.examples-row {
  display: grid !important;
  grid-template-columns: 1fr 1fr 1fr !important;
  gap: 10px !important;
  margin-top: 8px !important;
}
.example-btn button {
  background: var(--paper) !important;
  border: 1px solid var(--hair) !important;
  border-radius: 8px !important;
  padding: 12px 14px !important;
  font-family: 'Inter', sans-serif !important;
  font-size: 12.5px !important;
  font-weight: 500 !important;
  color: var(--ink) !important;
  text-align: left !important;
  cursor: pointer !important;
  width: 100% !important;
  min-height: 48px !important;
  line-height: 1.4 !important;
  white-space: normal !important;
  transition: all 0.15s ease !important;
}
.example-btn button:hover {
  background: var(--tint) !important;
  border-color: var(--rule) !important;
}

/* ── COMPARISON OUTPUT (TWO COLUMNS) ───────────────────── */
#comparison-row { gap: 18px !important; }

.lane {
  background: var(--paper);
  border: 1px solid var(--hair);
  border-radius: 10px;
  padding: 0;
  overflow: hidden;
  min-height: 200px;
}
.lane.left-lane { border-top: 4px solid var(--left-accent); }
.lane.right-lane { border-top: 4px solid var(--right-accent); }

.lane-header {
  padding: 16px 22px 12px 22px;
  border-bottom: 1px solid var(--hair);
  background: var(--off);
}
.lane-header .lane-title {
  font-family: 'Inter', sans-serif;
  font-size: 16px;
  font-weight: 800;
  color: var(--ink);
  letter-spacing: -0.3px;
  margin: 0;
}
.lane-header .lane-meta {
  font-family: 'JetBrains Mono', monospace;
  font-size: 11px;
  color: var(--muted);
  margin-top: 4px;
}

.lane-body, .lane-body * {
  font-family: 'Inter', sans-serif !important;
  font-size: 13.5px !important;
  line-height: 1.65 !important;
  color: var(--ink) !important;
}
.lane-body {
  padding: 20px 24px 24px 24px !important;
}
.lane-body strong { font-weight: 700 !important; }
.lane-body em { color: var(--muted) !important; font-style: italic !important; }
.lane-body p { margin: 0 0 10px 0 !important; }
.lane-body ol, .lane-body ul { padding-left: 22px !important; margin: 6px 0 10px 0 !important; }
.lane-body li { margin-bottom: 4px !important; }

/* ── FOOTER ────────────────────────────────────────────── */
#footer-block {
  text-align: center;
  margin-top: 48px;
  padding-top: 24px;
  border-top: 1px solid var(--hair);
  font-size: 12px;
  color: var(--muted);
  line-height: 1.6;
}
#footer-block a {
  color: var(--ink);
  text-decoration: underline;
  text-underline-offset: 2px;
}
"""


HEADER_HTML = """
<div id="header-block">
  <div class="title-row">Demo 03 · Stay Curious · Fine-Tuning LLMs</div>
  <h1>Junior Associate</h1>
  <div class="subtitle">
    A small open-source model fine-tuned to do <b>first-pass contract review</b>
    in the house style of a fictional B2B law firm.
  </div>

  <div class="use-case">
    <div class="use-case-row">
      <div class="use-case-cell">
        <div class="use-case-label">THE FIRM</div>
        <div class="use-case-text">
          <b>Lexwell Advisors</b> (fictional) reviews contracts for SMBs.
          Every reply must follow the firm&rsquo;s exact format.
        </div>
      </div>
      <div class="use-case-cell">
        <div class="use-case-label">THE HOUSE STYLE</div>
        <div class="use-case-text">
          Top disclaimer &rarr; <b>I</b>ssue &middot; <b>R</b>ule &middot; <b>A</b>pplication &middot; <b>C</b>onclusion
          &rarr; numbered redlines &rarr; bottom disclaimer &rarr;
          <i>&mdash; Lexwell Advisors</i>
        </div>
      </div>
      <div class="use-case-cell">
        <div class="use-case-label">WHY FINE-TUNE</div>
        <div class="use-case-text">
          Instead of a 2,000-token system prompt on every call,
          we fine-tuned Qwen-3B on <b>80 example memos</b>.
          The structure now lives in a 30&nbsp;MB adapter.
        </div>
      </div>
    </div>
  </div>

  <div class="how-to-read">
    <b>How to read this:</b> paste a contract question below.
    The <b>left column</b> shows the base model with no fine-tuning &mdash; helpful, but unstructured.
    The <b>right column</b> shows the same model with the LoRA adapter &mdash; locked to Lexwell&rsquo;s format every time.
  </div>

  <div class="meta-pill">
    <span class="label">MODEL</span>
    <a href="https://huggingface.co/Curious-PM/lexwell-contract-irac-qwen2.5-3b-lora" target="_blank">Curious-PM/lexwell-contract-irac-qwen2.5-3b-lora</a>
    <span style="color: var(--rule);">·</span>
    <span style="color: var(--muted);">trained via</span>
    <a href="https://huggingface.co/blog/hf-skills-training" target="_blank">hf-llm-trainer</a>
  </div>
</div>
"""

FOOTER_HTML = """
<div id="footer-block">
  Built for <a href="https://curious.pm" target="_blank">Curious PM</a> &middot; Stay Curious session on fine-tuning &middot; <a href="https://huggingface.co/Curious-PM/lexwell-contract-irac-qwen2.5-3b-lora" target="_blank">View the model</a>
</div>
"""

PLACEHOLDER = "_The reply will appear here. First request takes ~15s while the GPU warms up._"


with gr.Blocks(title="Junior Associate · Stay Curious", css=CSS, theme=gr.themes.Base()) as demo:

    gr.HTML(HEADER_HTML)

    gr.HTML('<div class="section-label">Ask a contract question</div>')

    question = gr.Textbox(
        placeholder="e.g. Our SaaS vendor wants perpetual ML training rights on our customer data. Is that normal?",
        lines=3,
        show_label=False,
        container=False,
    )

    submit = gr.Button("Compare base vs fine-tuned →", elem_classes="primary-button", variant="primary")

    gr.HTML('<div class="section-label">Or try one of these</div>')

    with gr.Row(elem_classes="examples-row"):
        example_buttons = []
        for label, full_text in zip(EXAMPLE_LABELS, EXAMPLES):
            btn = gr.Button(label, elem_classes="example-btn")
            example_buttons.append((btn, full_text))

    with gr.Row(elem_id="comparison-row"):
        with gr.Column(elem_classes="lane left-lane"):
            gr.HTML(
                '<div class="lane-header">'
                '  <div class="lane-title">Base Qwen 2.5-3B</div>'
                '  <div class="lane-meta">no fine-tuning</div>'
                '</div>'
            )
            base_output = gr.Markdown(value=PLACEHOLDER, elem_classes="lane-body")

        with gr.Column(elem_classes="lane right-lane"):
            gr.HTML(
                '<div class="lane-header">'
                '  <div class="lane-title">+ LoRA adapter (fine-tuned)</div>'
                '  <div class="lane-meta">trained on 80 contract-review memos</div>'
                '</div>'
            )
            ft_output = gr.Markdown(value=PLACEHOLDER, elem_classes="lane-body")

    gr.HTML(FOOTER_HTML)

    # Wire up
    submit.click(fn=generate_both, inputs=question, outputs=[base_output, ft_output])
    question.submit(fn=generate_both, inputs=question, outputs=[base_output, ft_output])

    for btn, full_text in example_buttons:
        btn.click(fn=lambda t=full_text: t, inputs=None, outputs=question).then(
            fn=generate_both, inputs=question, outputs=[base_output, ft_output]
        )


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