#!/usr/bin/env python3
"""FR-Start — Fahrenheit Research incorporation advisor CLI (fully local, no API).

Runs on Ollama. Usage:
    python fr_start.py            interactive chat
    python fr_start.py "question" one-shot answer
"""
import sys
from pathlib import Path

import ollama

ROOT = Path(__file__).parent
# ponytail: gemma3:12b — strongest model already on this machine; drop to qwen2.5:7b-instruct if too slow
MODEL = "gemma3:12b"
NUM_CTX = 32768  # system prompt is ~13k tokens; default ctx would silently truncate the corpus


# ponytail: whole corpus in context (~30KB), add retrieval when corpus > ~200KB
def build_system() -> str:
    prompt = (ROOT / "system_prompt.md").read_text()
    corpus = "\n\n---\n\n".join(
        f.read_text() for f in sorted((ROOT / "corpus").glob("*.md"))
    )
    # scope rule repeated after the corpus — small models weight the end of long prompts
    tail = (
        "# Final rule (absolute)\n"
        "Decision procedure for every user message, in this order:\n"
        "1. Does it involve companies, incorporation, entities, taxes or tax rates, "
        "compliance, banking, payroll, founder visas, funding, grants, cross-border "
        "structures, or exits — in or between the US, India, UAE, Singapore, or UK? "
        "If YES: answer it from the corpus. This includes short factual questions like "
        "'What is Singapore's corporate tax rate?' or 'Which Dubai zone for fintech?'. "
        "NEVER give the refusal reply to these.\n"
        "2. Only if the message is clearly unrelated (code, poems, trivia, math, health, "
        "other countries, casual chat): give the standard FR-Start reply from the Scope "
        "section, nothing else. The standard reply is always the ENTIRE response — "
        "never append it before or after an answer, and never use it when you have "
        "answered the question.\n"
        "Formatting: simple markdown — bold key terms, '- ' lists, and small markdown "
        "tables for comparisons and the Decision card."
    )
    return f"{prompt}\n\n# Reference corpus\n\n{corpus}\n\n{tail}"


MARKER = "This is FR-Start"  # opening of the standard scope-refusal reply


def stream_reply(messages: list[dict]):
    """Stream the model's reply; cut off a scope-refusal wrongly appended after a real answer.

    ponytail: 12B model sometimes tacks the refusal boilerplate onto valid answers —
    enforcing in code beats another prompt nudge. Holds back a small tail so the
    marker can't slip through split across chunks.
    """
    pending = ""
    hold = len(MARKER) + 8
    for chunk in ollama.chat(model=MODEL, messages=messages, stream=True,
                             options={"num_ctx": NUM_CTX, "temperature": 0}):
        pending += chunk["message"]["content"]
        i = pending.find(MARKER)
        if i > 0 and pending[:i].strip():
            yield pending[:i].rstrip()  # real answer followed by boilerplate — truncate
            return
        if i != 0 and len(pending) > hold:
            yield pending[:-hold]
            pending = pending[-hold:]
    yield pending


def chat() -> None:
    system = build_system()
    messages: list[dict] = [{"role": "system", "content": system}]
    one_shot = sys.argv[1] if len(sys.argv) > 1 else None

    print(f"FR-Start (local: {MODEL}) — where should you incorporate? (US / India / UAE / Singapore / UK)")
    print("Ctrl-C or 'quit' to exit. First answer is slow while the model loads.\n")

    while True:
        if one_shot:
            user = one_shot
        else:
            try:
                user = input("you> ").strip()
            except (EOFError, KeyboardInterrupt):
                print()
                return
            if not user or user.lower() in ("quit", "exit"):
                return

        messages.append({"role": "user", "content": user})
        print()
        reply = ""
        for piece in stream_reply(messages):
            reply += piece
            print(piece, end="", flush=True)
        print("\n")
        messages.append({"role": "assistant", "content": reply})

        if one_shot:
            return


if __name__ == "__main__":
    chat()