import os
import json
from openai import AsyncOpenAI
from sqlmodel import Session, select
from fastapi.responses import StreamingResponse
from uuid import UUID

import models

# Point the standard AsyncOpenAI client to your local Ollama server
client = AsyncOpenAI(
    base_url=os.getenv("LLM_BASE_URL", "http://host.docker.internal:11434/v1"),
    api_key=os.getenv("OPENAI_API_KEY", "ollama"), # Required by the SDK, but Ollama ignores it
)

async def stream_tutor_response(user_id: UUID, user_message: str, session: Session):
    # 1. Fetch Context: Get the student's active weak areas
    active_mistakes = session.exec(
        select(models.Mistake)
        .where(models.Mistake.user_id == user_id)
        .where(models.Mistake.is_mastered == False)
    ).all()
    
    weak_topics = list(set([m.topic for m in active_mistakes]))
    context_string = ", ".join(weak_topics) if weak_topics else "General Syllabus"

    # 2. Construct the highly-specific System Prompt
    system_prompt = f"""
    You are PragatiVeer's elite AI Tutor for JEE & NEET aspirants. 
    Your goal is to guide students to the answer, not just give it to them.
    
    STUDENT CONTEXT:
    This student currently struggles with the following topics: {context_string}.
    Keep this in mind and relate explanations back to foundational concepts in these areas if applicable.
    
    RULES:
    1. Respond in the student's preferred language (English or Hindi).
    2. Use LaTeX wrapped in single $ for inline math and $$ for block equations.
    3. Keep responses concise and focused on problem-solving steps.
    """

    # 3. Call the LLM
    response = await client.chat.completions.create(
        model=os.getenv("LLM_MODEL_NAME", "gemma"),
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message}
        ],
        stream=True,
    )

    # 4. Yield tokens as they arrive
    async for chunk in response:
        if chunk.choices[0].delta.content is not None:
            yield chunk.choices[0].delta.content