agent.py

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple

from llm import LLMError, chat_completion
from tools import (
    Message,
    StudentNotFoundError,
    end_exam,
    get_last_session_id,
    get_next_topic,
    set_current_session,
    start_exam,
)

EMAIL_RE = re.compile(r"([a-zA-Z0-9_.+\-]+@[a-zA-Z0-9\-]+\.[a-zA-Z0-9\-.]+)")


def _utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _looks_like_idk(text: str) -> bool:
    t = text.strip().lower()
    triggers = ["не знаю", "не пам'ятаю", "не памятаю", "не можу", "no idea", "idk", "не впевнений", "не впевнена"]
    return any(x in t for x in triggers)


@dataclass
class ExamAgent:
    stage: str = "collect_identity"  
    name: Optional[str] = None
    email: Optional[str] = None

    session_id: Optional[str] = None
    topics_total: int = 0
    current_topic: Optional[str] = None

    max_questions_per_topic: int = 3
    questions_in_topic: int = 0

    topic_scores: Dict[str, float] = field(default_factory=dict)
    history: List[Message] = field(default_factory=list)

    def _log(self, role: str, content: str) -> None:
        self.history.append(Message(role=role, content=content, datetime=_utc_now_iso()))

    def initial_message(self) -> str:
        msg = "Привіт! Я екзаменатор. Як тебе звати?"
        self._log("system", msg)
        return msg

    def _bind_session(self) -> None:
        if self.session_id:
            set_current_session(self.session_id)

    def _start_exam_tools(self) -> Tuple[bool, str]:
        self._log("tool_call", f"start_exam(email={self.email}, name={self.name})")
        try:
            topics = start_exam(self.email or "", self.name or "")
        except StudentNotFoundError:
            self.session_id = None
            self.topics_total = 0
            self.current_topic = None
            self.stage = "collect_identity"
            msg = (
                "Я не знайшов(ла) студента з таким email у списку. "
                "Перевір, будь ласка, email і надішли його ще раз."
            )
            self._log("system", msg)
            return False, msg

        self.session_id = get_last_session_id()
        self.topics_total = len(topics)

        msg = f"Добре, {self.name}. Починаємо іспит. Тем буде {self.topics_total}."
        self._log("system", msg)
        return True, msg

    def _next_topic(self) -> Optional[str]:
        self._bind_session()
        topic = get_next_topic()
        if topic:
            self.current_topic = topic
            self.questions_in_topic = 0
            return topic
        self.current_topic = None
        return None

    def _ask_question(self, api_key: str, model: str, base_url: str) -> str:
        assert self.current_topic is not None

        sys = "Ти екзаменатор. Питай коротко українською. Одне питання за раз."
        user = f"Тема: {self.current_topic}\nЗадай наступне питання."
        try:
            q = chat_completion(
                api_key=api_key,
                model=model,
                base_url=base_url,
                messages=[{"role": "system", "content": sys}, {"role": "user", "content": user}],
            ).strip()
        except LLMError:
            q = f"Поясни ключові поняття та приклад по темі: {self.current_topic}."

        self.questions_in_topic += 1
        self._log("system", q)
        return q

    def _evaluate_answer(self, api_key: str, model: str, base_url: str, answer: str) -> Dict[str, object]:
        assert self.current_topic is not None

        if _looks_like_idk(answer):
            return {"score": 0.0, "action": "next_topic", "note": "student_idk", "feedback": ""}

        sys = "Оціни відповідь студента для однієї теми. Поверни ТІЛЬКИ JSON."
        user = {
            "topic": self.current_topic,
            "answer": answer,
            "json_schema": {"score": "0..10", "action": "ask_followup|next_topic", "feedback": "string"},
        }

        try:
            raw = chat_completion(
                api_key=api_key,
                model=model,
                base_url=base_url,
                messages=[{"role": "system", "content": sys}, {"role": "user", "content": json.dumps(user, ensure_ascii=False)}],
            ).strip()
            data = json.loads(raw)

            score = float(data.get("score", 0.0))
            score = max(0.0, min(10.0, score))

            action = str(data.get("action", "ask_followup")).strip()
            if action not in {"ask_followup", "next_topic"}:
                action = "ask_followup"

            feedback = str(data.get("feedback", "")).strip()
            return {"score": score, "action": action, "feedback": feedback}
        except Exception:
            length = len(answer.strip())
            score = 3.0 if length < 80 else 6.0 if length < 250 else 7.5
            action = "next_topic" if score >= 7.0 else "ask_followup"
            return {"score": score, "action": action, "feedback": "Оцінка приблизна (fallback)."}

    def _finalize(self) -> str:
        avg = round((sum(self.topic_scores.values()) / max(1, len(self.topic_scores))) if self.topic_scores else 0.0, 1)

        strong = [t for t, s in self.topic_scores.items() if s >= 7.0]
        weak = [t for t, s in self.topic_scores.items() if s < 7.0]

        feedback_lines = []
        if strong:
            feedback_lines.append("Добре вийшло по темах: " + ", ".join(strong) + ".")
        if weak:
            feedback_lines.append("Варто підтягнути: " + ", ".join(weak) + ".")
        if not feedback_lines:
            feedback_lines.append("Дякую! Є над чим попрацювати — продовжуй практикуватися.")

        msg = f"Іспит завершено. Оцінка: {avg}/10.\n" + "\n".join(feedback_lines)
        self._log("system", msg)

        self._log("tool_call", f"end_exam(email={self.email}, score={avg}, history=[...])")
        end_exam(self.email or "", avg, [m.to_dict() for m in self.history])

        self.stage = "finished"
        return msg

    def step(self, user_text: str, api_key: str, model: str, base_url: str) -> str:
        user_text = (user_text or "").strip()
        self._log("user", user_text)

        if self.stage == "collect_identity":
            if not self.name:
                self.name = user_text if user_text else None
                if not self.name:
                    msg = "Як тебе звати?"
                    self._log("system", msg)
                    return msg
                msg = f"Приємно познайомитись, {self.name}! Тепер напиши свій email."
                self._log("system", msg)
                return msg

            if not self.email:
                m = EMAIL_RE.search(user_text)
                self.email = m.group(1) if m else None
                if not self.email:
                    msg = "Не бачу коректного email. Спробуй ще раз (наприклад: name@example.com)."
                    self._log("system", msg)
                    return msg

                ok, start_msg = self._start_exam_tools()
                if not ok:
                    self.email = None
                    return start_msg

                self.stage = "exam"
                topic = self._next_topic()
                if not topic:
                    return self._finalize()

                intro = f"Тема 1/{self.topics_total}: {topic}"
                self._log("system", intro)
                q = self._ask_question(api_key, model, base_url)
                return intro + "\n" + q

        if self.stage == "exam":
            if not self.current_topic:
                return self._finalize()

            eval_res = self._evaluate_answer(api_key, model, base_url, user_text)
            score = float(eval_res.get("score", 0.0))
            action = str(eval_res.get("action", "ask_followup"))
            feedback = str(eval_res.get("feedback", "")).strip()

            prev = self.topic_scores.get(self.current_topic, 0.0)
            self.topic_scores[self.current_topic] = max(prev, score)

            idk = (eval_res.get("note") == "student_idk")
            too_many = self.questions_in_topic >= self.max_questions_per_topic
            good_enough = score >= 7.0

            if idk or good_enough or too_many or action == "next_topic":
                note = f"Коментар: {feedback}\n" if feedback else ""
                next_t = self._next_topic()
                if not next_t:
                    return note + self._finalize()

                done = len({t for t in self.topic_scores.keys()})
                header = f"{note}Переходимо далі.\nТема {min(done+1, self.topics_total)}/{self.topics_total}: {next_t}"
                self._log("system", header)
                q = self._ask_question(api_key, model, base_url)
                return header + "\n" + q

            follow = "Ок. Уточню:"
            self._log("system", follow)
            q = self._ask_question(api_key, model, base_url)
            if feedback:
                return f"Коментар: {feedback}\n{follow}\n{q}"
            return f"{follow}\n{q}"

        msg = "Іспит уже завершено. Якщо хочеш — натисни Reset і почнемо заново."
        self._log("system", msg)
        return msg