# generation_service_anthropic.py
# For LangSmith tracing; NO Braintrust; clean for OpenAI/Anthropic API

import os
import anthropic
import re
import traceback
from typing import List, Dict, AsyncGenerator
from langsmith import traceable

# --- Environment: ensure API key is injected (from Replit secrets) ---
os.environ["LANGSMITH_ENDPOINT"] = "https://api.smith.langchain.com"
os.environ["LANGSMITH_TRACING"] = "true"
os.environ["ANTHROPIC_API_KEY"] = os.environ["ANTHROPIC_API_KEY"]
os.environ["LANGSMITH_API_KEY"] = os.environ["LANGSMITH_API_KEY"]
os.environ["LANGSMITH_PROJECT"] = os.environ["LANGSMITH_PROJECT"]

# --- Anthropic config and client ---
ANTHROPIC_API_KEY = os.environ["ANTHROPIC_API_KEY"]
GENERATION_MODEL = "claude-3-7-sonnet-20250219"
client = anthropic.AsyncAnthropic(api_key=ANTHROPIC_API_KEY) if ANTHROPIC_API_KEY else None

def check_generator_status():
    if not client: return False, "Anthropic client not initialized."
    return True, f"Anthropic generation service ready (Model: {GENERATION_MODEL})."

def clean_source_text(text):
    if not text: return ""
    cleaned = text; cleaned = re.sub(r'@\d+', '', cleaned); cleaned = re.sub(r'<HAL>', '', cleaned, flags=re.IGNORECASE); cleaned = cleaned.replace('<br>', ' ').replace('<br />', ' '); cleaned = re.sub(r'\s+', ' ', cleaned).strip()
    return cleaned

def format_context_for_prompt(documents):
    if not documents: return ""
    formatted_docs = []
    language_key = 'hebrew_text'; id_key = 'original_id'
    for index, doc in enumerate(documents):
        full_text_original = doc.get(language_key, ''); doc_id = doc.get(id_key, f'unknown_{index+1}')
        full_text_cleaned = clean_source_text(full_text_original)
        if full_text_cleaned: formatted_docs.append(f"<source index=\"{index + 1}\" id=\"{doc_id}\">\n<full_text>{full_text_cleaned}</full_text>\n</source>")
    return "\n\n".join(formatted_docs)

EXAMPLE_RESPONSE_HEBREW = """<example response hebrew>
על פי המקורות שהובאו, חשיבות השמחה בעבודת ה' היא מרכזית. נאמר כי <quote source_index="1">עיקר עבודת ה' היא בשמחה</quote>, כפי הפסוק <quote source_index="1">'עבדו את ה' בשמחה'</quote>. הסיבה לכך היא <quote source_index="1">כי השמחה פותחת הלב ומאירה הנשמה, ומביאה לידי דביקות בהשי"ת</quote>. לעומת זאת, מצב של עצבות גורם לתוצאה הפוכה, שכן <quote source_index="1">על ידי העצבות ח"ו נסתם הלב ואינו יכול לקבל אור הקדושה</quote>. מקור נוסף מדגיש כי השמחה היא תנאי לקבלת רוח הקודש והשראת השכינה, כפי שנאמר <quote source_index="2">שאין השכינה שורה אלא מתוך שמחה של מצוה</quote>, וכן <quote source_index="2">שלא שרתה עליו שכינה מפני שהיה עצב</quote>, כפי שלמדו מיעקב אבינו.
</example response hebrew>"""

@traceable
async def generate_response_stream_async(
    messages: List[Dict],
    context_documents: List[Dict],
) -> AsyncGenerator:
    """
    Generates a response using Anthropic, yields text chunks.
    Traced with LangSmith.
    """
    global client
    ready, msg = check_generator_status()
    if not ready or client is None: yield f"--- שגיאה: {msg} ---"; return

    last_user_msg_content = "שאלה לא נמצאה"
    for msg_ in reversed(messages):
        if msg_.get("role") == "user": last_user_msg_content = str(msg_.get("content", "")); break

    try:
        formatted_context = format_context_for_prompt(context_documents)
        has_context = bool(formatted_context)
        if not has_context and context_documents:
            yield f"--- שגיאה: המקורות שסופקו ריקים לאחר ניקוי. ---"; return
        elif not has_context and not context_documents:
            yield f"--- שגיאה: לא סופקו מקורות להקשר. ---"; return
    except Exception as format_err:
        yield f"--- שגיאה בעיצוב ההקשר: {format_err} ---"; return

    # System prompt as before
    system_prompt = f"""<instructions>
You are an expert assistant specializing in Chassidic texts...
**Response Requirements:**
(Keep all instructions as before)
</instructions>

{EXAMPLE_RESPONSE_HEBREW}"""

    api_messages = []
    user_prompt_content = f"<context>\n{formatted_context}\n</context>\n\nBased *exclusively* on the source text provided... Question (Hebrew):\n{last_user_msg_content}"
    api_messages.append({"role": "user", "content": user_prompt_content})

    print(f" -> Sending request to Anthropic (Model: {GENERATION_MODEL})...")
    final_response_text_chunks = []

    try:
        async with client.messages.stream(
            model=GENERATION_MODEL, max_tokens=20000, system=system_prompt,
            messages=api_messages, temperature=1.0,
            thinking={"type": "enabled", "budget_tokens": 16000}
        ) as stream:
            print(f" -> Anthropic stream created successfully...")
            async for chunk in stream.text_stream:
                if chunk and chunk.strip():
                    final_response_text_chunks.append(chunk)
                    yield chunk

    except Exception as e:
        yield f"\n\n--- שגיאה: {type(e).__name__} - {e} ---"
        traceback.print_exc()