backend/main.py

import os
import sys
import google.generativeai as genai
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Dict, Any, Union, Optional
from dotenv import load_dotenv

# Import local modules
from . import rag
from . import tools

# 1. Load Environment Variables
load_dotenv()
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

if not GEMINI_API_KEY:
    print("⚠️  WARNING: GEMINI_API_KEY not found. Check your .env or Cloud Dashboard.")

genai.configure(api_key=GEMINI_API_KEY)


MODEL_NAME = 'gemini-2.5-flash'


# --- API Models ---

class ChatRequest(BaseModel):
    message: str


class ChatResponse(BaseModel):
    response: str
    tool_code: Optional[str] = None
    tool_data: Optional[Union[Dict[str, Any], List[Any]]] = None


# --- FastAPI App Setup ---
app = FastAPI(
    title="Tharushika's AI Portfolio API",
    description="Backend for AI Portfolio using Gemini & RAG",
    version="1.0.0"
)

# --- CORS MIDDLEWARE ---
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("startup")
async def startup():
    """Initialize the RAG system when the server boots up."""
    rag.initialize_rag()

# --- Endpoints ---


@app.get("/")
def health_check():
    """Simple check to see if the API is running."""
    return {"status": "ok", "message": "Portfolio API is live"}


# --- CRITICAL FIX: Removed 'async' below ---
# This allows FastAPI to run this blocking code in a separate thread,
# keeping the server responsive (fixing the /docs loading issue).
@app.post("/chat", response_model=ChatResponse)
def chat_endpoint(request: ChatRequest):
    user_msg = request.message

    # --- SYSTEM INSTRUCTION ---
    # This guides the model to be a conversationalist first, and a tool-user second.
    sys_instruct = """
    You are the AI portfolio assistant for Tharushika Abedheera.
    
    GUIDELINES:
    1. Do NOT use tools (like get_projects or get_skills) for general questions like "Tell me about yourself", "Who are you?", "Hi", or "What do you do?".
    2. For general questions, answer textually using your internal knowledge or the provided RAG context.
    3. ONLY call 'get_projects' if the user specifically asks to SEE a list of projects, code, or portfolio items.
    4. ONLY call 'get_skills' if the user specifically asks to SEE the tech stack or skills list.
    """

    # 1. Initialize Gemini with Toolkit AND System Instruction
    model = genai.GenerativeModel(
        MODEL_NAME,
        tools=tools.ALL_TOOLS_LIST,
        system_instruction=sys_instruct
    )

    try:
        # 2. Ask Gemini the question
        chat = model.start_chat(enable_automatic_function_calling=False)
        response = chat.send_message(user_msg)

        # 3. Check for Function Calls (The "Router")
        function_call = None
        if response.parts:
            for part in response.parts:
                if part.function_call:
                    function_call = part.function_call
                    break

        # --- PATH A: Tool Triggered (Structured Data) ---
        if function_call:
            tool_name = function_call.name
            print(f"🛠️  Tool Triggered: {tool_name}")

            if tool_name in tools.TOOL_FUNCTIONS:
                # A. Execute the Python function
                data = tools.TOOL_FUNCTIONS[tool_name]()

                # B. Get the frontend signal code
                code = tools.TOOL_CODE_MAP.get(tool_name)

                # C. Get RAG Context (The Summary)
                # We retrieve context based on the user's message (e.g., "Show me projects")
                context_chunks = rag.retrieve_context(user_msg, k=2)
                context_text = "\n".join(context_chunks)

                # D. Generate the Intro using the Context
                # --- UPDATED PROMPT HERE ---
                intro_prompt = f"""
                The user asked: '{user_msg}'. 
                You just triggered the tool '{tool_name}' to show them visual data.
                
                Your Task:
                Write a professional and engaging introduction (2-3 sentences) for this data.
                
                CRITICAL INSTRUCTION:
                The text below contains a specific summary written by the owner about this section.
                YOU MUST incorporate this context into your answer to explain the user's specific experience in this area.
                
                CONTEXT (USE THIS): 
                {context_text}
                """

                intro_model = genai.GenerativeModel(MODEL_NAME)
                intro_resp = intro_model.generate_content(intro_prompt)
                intro_text = intro_resp.text.strip()

                return ChatResponse(
                    response=intro_text,
                    tool_code=code,
                    tool_data=data
                )

        # --- PATH B: Pure RAG Chat (Unstructured Context) ---
        print("🧠  RAG Path Triggered")

        # 1. Retrieve relevant text chunks
        context_chunks = rag.retrieve_context(user_msg, k=4)
        context_text = "\n\n".join(context_chunks)

        # 2. Construct the Prompt with a FALLBACK BIO
        rag_prompt = f"""
        You are an AI assistant for Tharushika Abedheera's portfolio. He is a Machine Learning Engineer & AI Specialist. And yes he made this RAG AI based protfolio as well.
        Your goal is to answer the user's question professionally and confidently, acting as Tharushika.

        STRICT RULES:
        - Use ONLY the context provided below.
        - If the answer isn't in the context, use the FALLBACK BIO information.
        - Keep answers concise (under 3-4 sentences) unless asked for detail.

        CONTEXT FROM KNOWLEDGE BASE:
        {context_text}

        FALLBACK BIO (Use this if Context is empty or insufficient):
        Tharushika Abedheera Machine Learning Engineer with a research-driven approach, specializing in NLP, Generative AI, and AI agent frameworks. Passionate about transforming AI/ML research into real-world applications, with hands-on experience and expertise in Machine learning, Deep learning, large language models (LLMs), retrieval-augmented generation (RAG), and autonomous AI systems. Adept at designing, deploying, and optimizing AI-driven systems to enhance scalability and innovation..

        USER QUESTION:
        {user_msg}
        """

        # 3. Generate Answer
        rag_model = genai.GenerativeModel(MODEL_NAME)
        text_response = rag_model.generate_content(rag_prompt).text

        return ChatResponse(response=text_response, tool_code=None, tool_data=None)

    except Exception as e:
        print(f"❌ Error handling chat: {e}")
        raise HTTPException(status_code=500, detail="Internal Server Error")