# File: src/interview_agent.py
# Purpose: Stateful LLM interview agent using Groq (llama-3.3-70b-versatile)

from typing import List, Dict, Optional
from groq import Groq

from config import (
    GROQ_API_KEY, GROQ_LLM_MODEL, LLM_TEMPERATURE, LLM_MAX_TOKENS,
    MAX_QUESTIONS, INTRO_MESSAGE
)
from src.resume_rag import retrieve_resume_context

_groq = Groq(api_key=GROQ_API_KEY)

SYSTEM_PROMPT = """You are an expert AI technical recruiter conducting a structured phone interview.
Your job is to:
1. Ask one clear question at a time — never ask multiple questions together.
2. Listen carefully to the candidate's answer, identify skills and technical depth.
3. Ask a natural follow-up if the answer is vague or incomplete.
4. Move to the next planned question when the current topic is sufficiently explored.
5. Keep a professional, warm, and encouraging tone.

Rules:
- Never reveal you are evaluating the candidate.
- Do not provide hints or correct the candidate.
- Keep responses under 3 sentences.
- If the candidate is off-topic, gently redirect.
"""

BASE_QUESTIONS = [
    "Can you start by giving me a brief overview of your background and what you've been working on recently?",
    "What is the most technically challenging project you have worked on, and what was your specific contribution?",
    "How do you approach debugging a model that is underperforming on a validation set?",
    "Describe your experience with deploying machine learning models to production.",
    "How do you handle class imbalance in a classification problem?",
    "What is your experience with version control and collaborative development?",
    "Where do you see yourself growing technically in the next year?",
]


class InterviewAgent:
    def __init__(
        self,
        session_id: int,
        candidate_name: str,
        resume_skills: List[str],
        use_rag: bool = True,
    ):
        self.session_id = session_id
        self.candidate_name = candidate_name
        self.resume_skills = resume_skills
        self.use_rag = use_rag
        self.history: List[Dict[str, str]] = []
        self.question_index = 0
        self.turn_count = 0
        self.is_complete = False
        self._build_question_plan()

    def _build_question_plan(self):
        personalised = []
        for skill in self.resume_skills[:3]:
            personalised.append(
                f"I noticed your experience with {skill}. Can you walk me through a specific project where you used it?"
            )
        self.questions = (personalised + BASE_QUESTIONS)[:MAX_QUESTIONS]

    def _get_resume_context(self, last_answer: str) -> str:
        if not self.use_rag:
            return ""
        return retrieve_resume_context(self.session_id, last_answer)

    def get_next_ai_message(self, candidate_answer: Optional[str] = None) -> str:
        if candidate_answer is not None:
            self.history.append({"role": "user", "content": candidate_answer})
            self.turn_count += 1

        if self.question_index >= len(self.questions):
            self.is_complete = True
            closing = (
                f"Thank you so much, {self.candidate_name}. That concludes our screening interview. "
                "The hiring team will review your responses and reach out with next steps. "
                "Have a great day!"
            )
            self.history.append({"role": "assistant", "content": closing})
            return closing

        resume_ctx = ""
        if candidate_answer:
            resume_ctx = self._get_resume_context(candidate_answer)

        next_question = self.questions[self.question_index]
        system_with_context = SYSTEM_PROMPT
        if resume_ctx:
            system_with_context += f"\n\nRelevant resume context:\n{resume_ctx}"
        system_with_context += f"\n\nNext planned question: {next_question}"
        system_with_context += (
            "\n\nIf the last answer was incomplete, ask a short follow-up first. "
            "Otherwise transition naturally and ask the next planned question."
        )

        messages = [{"role": "system", "content": system_with_context}] + self.history

        response = _groq.chat.completions.create(
            model=GROQ_LLM_MODEL,
            messages=messages,
            temperature=LLM_TEMPERATURE,
            max_tokens=LLM_MAX_TOKENS,
        )
        ai_text = response.choices[0].message.content.strip()

        if candidate_answer is not None:
            self.question_index += 1

        self.history.append({"role": "assistant", "content": ai_text})
        return ai_text

    def get_full_transcript(self) -> List[Dict[str, str]]:
        transcript = []
        for msg in self.history:
            speaker = "AI" if msg["role"] == "assistant" else "Candidate"
            transcript.append({"speaker": speaker, "text": msg["content"]})
        return transcript