from fastapi import FastAPI
from pydantic import BaseModel
from retriever import retrieve
from memory import get_context
from llm import generate_response

app = FastAPI()


class Message(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    messages: list[Message]


@app.get("/health")
def health():

    return {
        "status": "ok"
    }


@app.post("/chat")
def chat(req: ChatRequest):

    try:

        # Get conversation history
        history = get_context(req.messages)

        # Latest user query
        query = req.messages[-1].content


        # Handle vague queries
        if len(query.split()) < 3:

            return {

                "reply":
                "Could you provide more details such as role, required skills, or experience level?",

                "recommendations": [],

                "end_of_conversation": False
            }


        # Retrieve assessments
        recommendations = retrieve(
            query,
            k=10
        )


        # Remove noisy results
        filtered = []

        bad_keywords = [

            "job control",
            "numerical reasoning"

        ]

        for r in recommendations:

            name = r["name"].lower()

            if not any(
                bad in name
                for bad in bad_keywords
            ):

                filtered.append(r)


        recommendations = filtered[:5]


        # Create readable text for LLM
        recommendation_text = ""

        for r in recommendations:

            recommendation_text += (

                f"- {r['name']} "
                f"(Type: {r['test_type']})\n"

            )


        prompt = f"""
You are an SHL assessment recommendation assistant.

Conversation:
{history}

Retrieved assessments:
{recommendation_text}

Rules:

1. Recommend ONLY assessments from the retrieved list.
2. Never invent new assessments.
3. Never infer abilities not explicitly mentioned.
4. Explain briefly why each assessment fits.
5. Do not output raw Python dictionaries.
6. Keep response under 120 words.
7. Ask follow-up questions if information is missing.
8. If query is unrelated to SHL assessments, politely refuse.
9. Use a professional conversational tone.
"""


        reply = generate_response(
            prompt
        )


        return {

            "reply": reply,

            "recommendations": recommendations,

            "end_of_conversation": False
        }


    except Exception as e:

        return {

            "reply":
            f"ERROR: {str(e)}",

            "recommendations": [],

            "end_of_conversation": False
        }