import os
from dataclasses import dataclass
from groq import Groq
from src.labdaps.config import GROQ_MODEL, MAX_HISTORY_MESSAGES
from src.labdaps.retrieval.retriever import RetrievedChunk, retrieve
from src.labdaps.ingestion.embedder import Embedder
from src.labdaps.chat.prompts import build_system_prompt


@dataclass
class ChatResponse:
    text: str
    sources: list[RetrievedChunk]


class ChatSession:
    def __init__(self, embedder: Embedder):
        self.embedder = embedder
        self.client = Groq(api_key=os.environ["GROQ_API_KEY"])
        self.history: list[dict] = []

    def reset(self):
        self.history = []

    def ask(self, question: str):
        chunks = retrieve(question, self.embedder)
        system_prompt = build_system_prompt(chunks)

        trimmed_history = self.history[-MAX_HISTORY_MESSAGES:]
        messages = (
            [{"role": "system", "content": system_prompt}]
            + trimmed_history
            + [{"role": "user", "content": question}]
        )

        full_text = ""
        stream = self.client.chat.completions.create(
            model=GROQ_MODEL,
            messages=messages,
            stream=True,
            max_tokens=2048,
        )
        for chunk in stream:
            delta = chunk.choices[0].delta.content or ""
            full_text += delta
            yield delta, chunks

        self.history.append({"role": "user", "content": question})
        self.history.append({"role": "assistant", "content": full_text})