import os
import sys
import torch

# Add workspace dir to Python path so we can import our updated app.py
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import app

def ask(query):
    context, url_to_title, prompt_note = app.retrieve_context(query)
    
    if context == "OUT_OF_SCOPE_REFUSAL":
        print("\n--- Assistant Response ---")
        print("I'm only authorized to assist with questions related to Mintoak Innovations Private Limited, its products, and services. How can I help you with Mintoak today? 😊")
        print("--------------------------\n")
        return

    # Prepare chat prompt matching our system configuration
    user_content = f"Context:\n{context}\n\nQuestion: {query}"
    if prompt_note:
        user_content += f"\n\nNote: {prompt_note}"
    messages = [
        {"role": "system", "content": app.SYSTEM_PROMPT},
        {"role": "user",   "content": user_content}
    ]
    prompt = app._tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = app._tokenizer([prompt], return_tensors="pt").to(app.device)
    
    print("\nThinking...", end="", flush=True)
    with torch.no_grad():
        outputs = app._model.generate(
            **inputs, 
            max_new_tokens=500, 
            temperature=0.1, 
            do_sample=False
        )
    response = app._tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
    cleaned_response = app.clean_response(response)
    
    # Append sources if not generated
    sources = app.build_sources(context, url_to_title)
    if "Source:" not in cleaned_response and sources:
        links = ", ".join(f"[{s['title']}]({s['url']})" for s in sources)
        cleaned_response += f"\n\n👉 For more details, visit: {links}"
        
    print("\r" + " " * 15 + "\r", end="") # Clear the 'Thinking...' line
    print("\n--- Assistant Response ---")
    print(cleaned_response)
    print("--------------------------\n")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python ask.py \"Your question here\"")
    else:
        query = " ".join(sys.argv[1:])
        ask(query)