from transformers import AutoTokenizer, AutoModelForCausalLM
#from user_data import load_user_data, save_user_data
from phonetics import analyze_audio_phonetically, extract_phonemes

model_name = "BeastGokul/Nika-1.5B"
llm_tokenizer = AutoTokenizer.from_pretrained(model_name)
llm_model = AutoModelForCausalLM.from_pretrained(model_name)

SYSTEM_PROMPT = """You are a specialized pronunciation assistant for non-native English speakers.\nYour job is to provide targeted, actionable feedback based on the user's speech or description.\n\nWhen analyzing pronunciation:\n1. Identify at most 2 specific phonemes or pronunciation patterns that need improvement\n2. Explain how the sound is correctly formed (tongue position, lip movement, etc.)\n3. Suggest one simple, targeted exercise for practice\n4. Be encouraging and note any improvements from previous sessions\n5. Use simple language appropriate for language learners\n\nWhen provided with phonetic analysis data, incorporate this information into your feedback.\n"""

def get_llm_feedback(audio=None, text=None, reference_text=None, user_id="default", transcribe_func=None):
    user_data = load_user_data(user_id)
    # Process audio if provided
    if audio:
        from user_data import save_audio
        audio_path = save_audio(audio, user_id)
        # Transcribe if no text was provided
        if not text and transcribe_func:
            text = transcribe_func(audio_path)
        # Get phonetic analysis
        phonetic_analysis = analyze_audio_phonetically(audio_path, reference_text)
        phonetic_info = f"""
Phonetic analysis:\n- Detected phonemes: {phonetic_analysis['detected_phonemes']}\n"""
        if reference_text:
            phonetic_info += f"- Reference phonemes: {phonetic_analysis.get('reference_phonemes', 'N/A')}\n"
    else:
        audio_path = None
        phonetic_info = ""
    # Get user history context
    history_context = ""
    if user_data["practice_sessions"]:
        phoneme_counts = {p: data["practice_count"] for p, data in user_data["phoneme_progress"].items()}
        challenging = sorted(phoneme_counts.items(), key=lambda x: x[1], reverse=True)[:3]
        history_context = f"""
User has practiced {len(user_data['practice_sessions'])} times before.\nCommon challenging phonemes: {', '.join([p for p, _ in challenging])}.\n"""
    # Build prompt for LLM
    if text:
        user_input = f"I said: '{text}'"
        if reference_text and reference_text != text:
            user_input += f". I was trying to say: '{reference_text}'"
    else:
        user_input = "Please analyze my pronunciation."
    full_prompt = f"""{SYSTEM_PROMPT}\n\nUser history:\n{history_context}\n\n{phonetic_info}\n\nUser: {user_input}\n"""
    # Get LLM response
    inputs = llm_tokenizer(full_prompt, return_tensors="pt").to(llm_model.device)
    import torch
    with torch.no_grad():
        outputs = llm_model.generate(
            **inputs, 
            max_new_tokens=200,
            temperature=0.7,
            top_p=0.9,
            do_sample=True
        )
    response = llm_tokenizer.decode(outputs[0], skip_special_tokens=True)
    try:
        response = response.split("Assistant: ")[-1].strip()
    except:
        pass
    # Track the session if audio was provided
    if audio_path:
        from user_data import track_practice_session
        track_practice_session(user_id, audio_path, text, reference_text, response)
    return response, text