#!/usr/bin/env python3
"""
Generate content pools for JingTian sample documents using Azure Claude API.
Run this once to create data/content_pools.json.
Makes multiple smaller API calls to avoid timeouts.
"""

import json
import os
import sys
import time
from pathlib import Path

try:
    from anthropic import AnthropicFoundry
except ImportError:
    print("Error: anthropic package not installed. Run: pip install anthropic")
    sys.exit(1)

ENDPOINT = "https://admin-ml6rv1i3-swedencentral.services.ai.azure.com/anthropic/"
MODEL = "claude-sonnet-4-6"
API_KEY = os.environ.get("AZURE_CLAUDE_API_KEY", "")

OUTPUT_PATH = Path(__file__).parent / "data" / "content_pools.json"

SYSTEM = """You are generating realistic content for sample documents used by a Hong Kong IP/trademark law firm called Jingtian & Gongcheng LLP. The sole practitioner is Benjamin Choi (蔡明睿), Partner. His office is at Suites 3203-3207, 32/F, Edinburgh Tower, The Landmark, 15 Queen's Road Central, Central, Hong Kong.

Return ONLY valid JSON arrays, no markdown code blocks, no explanation. Use placeholders like {deadline}, {tm_number}, {tm_text}, {client_name}, {contact_person}, {class} where indicated."""

PROMPTS = {
    "letter_to_client_bodies": """Generate a JSON array of 10 different letter bodies FROM Jingtian & Gongcheng TO clients regarding trademark matters. Each should be 2-4 paragraphs, formal HK legal style. Varied topics: opposition response needed, examination report received, renewal reminder, registration confirmed, filing update, amendment required, deadline approaching, evidence submission required, hearing notice, costs estimate.

Each MUST contain placeholders: {deadline}, {tm_number}, {tm_text}, {client_name}, {contact_person}.
Start each with "Dear {contact_person},".
Mix some bilingual (EN/CN) content naturally.""",
    "letter_from_client_bodies": """Generate a JSON array of 10 different letter bodies FROM clients TO Benjamin Choi at Jingtian. These are client instructions/responses, slightly less formal. Varied topics: proceed with filing, approve response, query about status, provide evidence, confirm renewal, change instructions, budget concerns, urgent request, new trademark idea, withdrawal request.

Use placeholders: {tm_number}, {tm_text} where relevant. About 70% should include {deadline}.
Start each with "Dear Mr. Choi," or "Dear Benjamin,".""",
    "email_from_client_bodies": """Generate a JSON array of 10 different SHORT email bodies FROM clients TO Benjamin Choi. Casual/professional email style, 1-3 short paragraphs each. Varied topics: quick follow-up, status check, forwarding a document, asking about costs, confirming a meeting, deadline reminder, new matter inquiry, sending signed docs, travel affecting timeline, board meeting deadline.

About 60% should include {deadline}. Use {tm_number} or {tm_text} where natural.
Start with "Hi Benjamin," or "Dear Benjamin," or "Hi Mr. Choi,".""",
    "memo_bodies": """Generate a JSON array of 8 different internal memo bodies. These are Benjamin Choi's internal notes/reminders about trademark matters. Each should list 2-3 items using numbered placeholders like {tm_number_1}, {tm_text_1}, {deadline_1}, {client_1}, {tm_number_2}, {tm_text_2}, {deadline_2}, {client_2}, etc.

Topics: upcoming renewals batch, overdue items, priority matters, quarterly review, opposition deadlines, examination response deadlines, new filings status, billing follow-up.""",
    "short_pools": """Generate a JSON object with these keys:

"email_subjects": array of 15 realistic email subject lines for HK trademark correspondence. Use {tm_number} or {tm_text} placeholders. Mix EN/CN.

"deadline_phrases": array of 15 different ways to express a deadline. Each contains {deadline}. Mix formal/informal, EN/CN. Examples: "Please respond by {deadline}", "The statutory deadline is {deadline}", "請於{deadline}前回覆".

"letter_closings": array of 8 formal letter closing paragraphs (before "Yours faithfully"). HK legal style.

"invoice_descriptions": array of 12 invoice line item descriptions for trademark services. Use {class} placeholder where relevant. e.g. "Professional fees for trademark application filing - Class {class}", "Government filing fee for trademark renewal".""",
}


def call_api(client, prompt, label):
    """Make a single API call and return parsed JSON."""
    print(f"  Generating {label}...", end=" ", flush=True)
    start = time.time()

    message = client.messages.create(
        model=MODEL,
        system=SYSTEM,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=8000,
    )

    raw = message.content[0].text.strip()

    # Strip markdown code blocks if present
    if raw.startswith("```"):
        lines = raw.split("\n")
        raw = "\n".join(lines[1:-1])

    elapsed = time.time() - start
    print(f"done ({elapsed:.1f}s, {len(raw)} chars)")
    return json.loads(raw)


def main():
    if not API_KEY:
        print("Error: AZURE_CLAUDE_API_KEY environment variable not set.")
        print("Usage: AZURE_CLAUDE_API_KEY=<key> python generate_content_pools.py")
        sys.exit(1)

    print(f"Connecting to Azure Claude ({MODEL})...")
    client = AnthropicFoundry(
        api_key=API_KEY,
        base_url=ENDPOINT,
    )

    content_pools = {}

    # Generate each pool separately
    for key, prompt in PROMPTS.items():
        try:
            result = call_api(client, prompt, key)

            if key == "short_pools":
                # This returns an object with multiple keys
                content_pools.update(result)
            else:
                content_pools[key] = result

        except json.JSONDecodeError as e:
            print(f"\n  ERROR parsing {key}: {e}")
            print(f"  Skipping {key}, continuing...")
            continue
        except Exception as e:
            print(f"\n  ERROR calling API for {key}: {e}")
            print(f"  Skipping {key}, continuing...")
            continue

    # Summary
    print("\nContent pools summary:")
    for key, val in content_pools.items():
        if isinstance(val, list):
            print(f"  {key}: {len(val)} items")

    # Save
    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    OUTPUT_PATH.write_text(
        json.dumps(content_pools, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    print(f"\nSaved to {OUTPUT_PATH}")


if __name__ == "__main__":
    main()