#!/usr/bin/env python3
"""
Tokenization bias test — local llama-server.
Compares EN vs ES token counts using your local model's tokenizer.

Usage: python3 test_local_tokenizer.py --host YOUR_HOST --port YOUR_PORT

Requires llama-server running with your model:
  llama-server -m your-model.gguf --port 8080

Part of: "Your AI charges you up to 67% more for not speaking English"
https://theprivatestack.com/research/tokenization-bias
"""

import argparse
import json
import urllib.request


def tokenize(text, base_url):
    url = f"{base_url}/tokenize"
    data = json.dumps({"content": text}).encode()
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req) as resp:
        result = json.loads(resp.read())
    return result["tokens"]


WORD_PAIRS = [
    ("contract", "contrato"),
    ("agreement", "acuerdo"),
    ("corporation", "corporación"),
    ("shareholders", "accionistas"),
    ("liability", "responsabilidad"),
    ("compliance", "cumplimiento"),
    ("notarization", "protocolización"),
    ("incorporation", "constitución"),
    ("bylaws", "estatutos"),
    ("board of directors", "consejo de administración"),
    ("power of attorney", "poder notarial"),
    ("due diligence", "debida diligencia"),
    ("articles of incorporation", "acta constitutiva"),
    ("tax identification number", "registro federal de contribuyentes"),
    ("artificial intelligence", "inteligencia artificial"),
    ("data sovereignty", "soberanía de datos"),
    ("retrieval augmented generation", "generación aumentada por recuperación"),
]

MX_TERMS = [
    "fideicomiso",
    "escritura pública",
    "acta constitutiva",
    "protocolo notarial",
    "Sociedad Anónima Promotora de Inversión de Capital Variable",
    "Registro Federal de Contribuyentes",
    "comprobantes fiscales digitales por Internet",
    "asamblea general ordinaria de accionistas",
    "Ley General de Sociedades Mercantiles",
    "Servicio de Administración Tributaria",
]

CLAUSES = {
    "NDA (EN)": "The Receiving Party agrees to hold in strict confidence all Confidential Information disclosed by the Disclosing Party and shall not disclose such information to any third party without prior written consent. This obligation shall survive the termination of this agreement for a period of five years.",
    "NDA (ES)": "La Parte Receptora se obliga a mantener en estricta confidencialidad toda la Información Confidencial revelada por la Parte Divulgante y no divulgará dicha información a terceros sin el consentimiento previo por escrito. Esta obligación sobrevivirá la terminación del presente contrato por un período de cinco años.",
    "Incorporation (EN)": "This Corporation is organized for the purpose of engaging in any lawful act or activity for which corporations may be organized under the General Corporation Law of the State of Delaware. The total number of shares of stock which the corporation shall have authority to issue is ten million shares of common stock, each having a par value of one cent per share.",
    "Constitutiva (ES)": "La Sociedad tiene por objeto la realización de cualquier acto o actividad lícita para la cual las sociedades pueden organizarse conforme a la Ley General de Sociedades Mercantiles del Estado Mexicano. El capital social autorizado será de diez millones de acciones ordinarias, cada una con valor nominal de un centavo por acción.",
    "Tax (EN)": "The taxpayer shall file annual returns with the Internal Revenue Service for each taxable year. All digital tax receipts must comply with current regulations and include the employer identification number assigned by the federal tax authority.",
    "Fiscal (ES)": "El contribuyente deberá presentar declaraciones anuales ante el Servicio de Administración Tributaria por cada ejercicio fiscal. Todos los comprobantes fiscales digitales por Internet deberán cumplir con la normatividad vigente e incluir el Registro Federal de Contribuyentes asignado por la autoridad fiscal federal.",
    "SAPI de CV (ES)": "PRIMERA.- DENOMINACIÓN, TIPO Y DURACIÓN. La Sociedad se denominará Ejemplo Tecnológico, Sociedad Anónima Promotora de Inversión de Capital Variable. La Sociedad tendrá una duración indefinida y se regirá por lo dispuesto en la Ley General de Sociedades Mercantiles, la Ley del Mercado de Valores y demás disposiciones legales aplicables.",
}

CLAUSE_PAIRS = [
    ("NDA (EN)", "NDA (ES)", "NDA"),
    ("Incorporation (EN)", "Constitutiva (ES)", "Constitutiva"),
    ("Tax (EN)", "Fiscal (ES)", "Fiscal"),
]


def main():
    parser = argparse.ArgumentParser(description="Test tokenizer bias against local llama-server")
    parser.add_argument("--host", default="localhost", help="llama-server host (default: localhost)")
    parser.add_argument("--port", default="8080", help="llama-server port (default: 8080)")
    args = parser.parse_args()

    base_url = f"http://{args.host}:{args.port}"

    try:
        tokenize("test", base_url)
        print(f"Connected to {base_url}\n")
    except Exception as e:
        print(f"Cannot connect to {base_url}: {e}")
        print("Make sure llama-server is running:")
        print("  llama-server -m your-model.gguf --port 8080")
        return

    # Word pairs
    print("=" * 90)
    print("  EN vs ES WORD PAIRS")
    print("=" * 90)
    print(f"{'English':<38} {'Tok':>4}  {'Spanish':<38} {'Tok':>4} {'D':>4}")
    print("-" * 90)

    total_en, total_es = 0, 0
    for en, es in WORD_PAIRS:
        t_en = len(tokenize(en, base_url))
        t_es = len(tokenize(es, base_url))
        total_en += t_en
        total_es += t_es
        print(f"{en:<38} {t_en:>4}  {es:<38} {t_es:>4} {t_es - t_en:>+4}")

    print("-" * 90)
    print(f"{'TOTAL':<38} {total_en:>4}  {'TOTAL':<38} {total_es:>4} {total_es - total_en:>+4}")
    print(f"\n  Word-level overhead: {((total_es / total_en) - 1) * 100:.1f}%\n")

    # Mexican terms
    print("=" * 90)
    print("  MEXICAN LEGAL/FISCAL TERMS")
    print("=" * 90)
    for term in MX_TERMS:
        tokens = tokenize(term, base_url)
        print(f'  "{term}" -> {len(tokens)} tokens')

    # Clauses
    print("\n" + "=" * 90)
    print("  CLAUSE COMPARISON")
    print("=" * 90)
    for label, text in CLAUSES.items():
        tokens = tokenize(text, base_url)
        print(f"  {label:<25} {len(text):>4} chars -> {len(tokens):>4} tokens ({len(text)/len(tokens):.1f} chars/tok)")

    print("\n  ES vs EN overhead:")
    for en_l, es_l, name in CLAUSE_PAIRS:
        en_t = len(tokenize(CLAUSES[en_l], base_url))
        es_t = len(tokenize(CLAUSES[es_l], base_url))
        print(f"    {name:<20} {((es_t / en_t) - 1) * 100:>+.1f}%")

    print()


if __name__ == "__main__":
    main()