backend/app/llm.py

import google.generativeai as genai
import json
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Validate required environment variable
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if not GEMINI_API_KEY:
    raise ValueError(
        "GEMINI_API_KEY environment variable is required. "
        "Please set it in your .env file or environment."
    )

# Configure Gemini
genai.configure(api_key=GEMINI_API_KEY)

# Import model router for multi-model rotation
from app.model_router import generate as router_generate, generate_with_info


async def generate_documentation(task_title: str, what_i_did: str, code_snippet: str | None = None) -> dict:
    """Generate docs for completed task. Returns {summary, details, tags}"""
    prompt = f"""
    Generate technical documentation for this completed work.

    Task: {task_title}
    What was done: {what_i_did}
    Code: {code_snippet or 'N/A'}

    Return ONLY valid JSON with:
    - "summary": one-line summary
    - "details": 2-3 paragraph technical documentation
    - "tags": array of 3-7 relevant tags

    Response must be pure JSON, no markdown.
    """

    # Use model router for multi-model rotation
    text = await router_generate(prompt, task_type="documentation")
    # Clean response (remove markdown code blocks if present)
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```")[1]
        if text.startswith("json"):
            text = text[4:]
    return json.loads(text.strip())


async def synthesize_answer(context: str, query: str) -> str:
    """Generate answer from context. Returns answer string."""
    prompt = f"""
    Based on this project memory:
    {context}

    Answer: {query}

    Cite specific entries. If info not found, say so.
    """

    # Use model router for multi-model rotation
    return await router_generate(prompt, task_type="synthesis")


async def get_embedding(text: str) -> list[float]:
    """Get embedding vector for text using Gemini embedding model."""
    result = genai.embed_content(
        model="models/text-embedding-004",
        content=text
    )
    return result['embedding']


async def generate_tasks(project_name: str, project_description: str, count: int = 50) -> list[dict]:
    """Generate demo tasks for a project using LLM.

    Args:
        project_name: Name of the project
        project_description: Description of the project
        count: Number of tasks to generate (max 50)

    Returns:
        List of tasks with title and description
    """
    # Cap at 50 tasks max
    count = min(count, 50)

    prompt = f"""
You are a project manager creating demo tasks for a hackathon project.

Project: {project_name}
Description: {project_description}

Generate exactly {count} simple, demo-friendly tasks for this software project. Each task should be:
- Simple and quick to complete (5-30 minutes each)
- Suitable for a demo or hackathon setting
- Cover typical software development activities (setup, coding, testing, docs, UI)

Include a mix of:
- Setup tasks (environment, dependencies, config)
- Feature implementation (simple features)
- Bug fixes (minor issues)
- Documentation (README, comments)
- Testing (basic tests)
- UI/UX improvements

Return ONLY a valid JSON array with objects containing:
- "title": short task title (max 100 chars)
- "description": brief description (1 sentence)

Example:
[
  {{"title": "Set up development environment", "description": "Install dependencies and configure local dev environment."}},
  {{"title": "Add user login button", "description": "Create a login button component in the header."}}
]

Return ONLY the JSON array, no markdown or extra text.
"""

    # Use model router for generation
    text = await router_generate(prompt, task_type="documentation")

    # Clean response (remove markdown code blocks if present)
    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        # Remove first and last lines (```json and ```)
        text = "\n".join(lines[1:-1])
        if text.startswith("json"):
            text = text[4:]

    try:
        tasks = json.loads(text.strip())
        # Validate structure
        if not isinstance(tasks, list):
            raise ValueError("Response is not a list")

        # Ensure each task has required fields
        validated_tasks = []
        for task in tasks:
            if isinstance(task, dict) and "title" in task:
                validated_tasks.append({
                    "title": str(task.get("title", ""))[:100],
                    "description": str(task.get("description", ""))
                })

        return validated_tasks
    except json.JSONDecodeError as e:
        raise ValueError(f"Failed to parse LLM response as JSON: {e}")


async def chat_with_tools(messages: list[dict], project_id: str) -> str:
    """Chat with AI using MCP tools for function calling.

    Args:
        messages: List of chat messages [{'role': 'user/assistant', 'content': '...'}]
        project_id: Project ID for context

    Returns:
        AI response string
    """
    from app.tools.projects import list_projects, create_project, join_project
    from app.tools.tasks import list_tasks, create_task, list_activity
    from app.tools.memory import complete_task, memory_search
    from app.model_router import router

    # Define tools for Gemini function calling
    tools = [
        {
            "name": "list_projects",
            "description": "List all projects for a user",
            "parameters": {
                "type": "object",
                "properties": {
                    "userId": {"type": "string", "description": "User ID"}
                },
                "required": ["userId"]
            }
        },
        {
            "name": "list_tasks",
            "description": "List all tasks for a project",
            "parameters": {
                "type": "object",
                "properties": {
                    "projectId": {"type": "string", "description": "Project ID"},
                    "status": {"type": "string", "enum": ["todo", "in_progress", "done"]}
                },
                "required": ["projectId"]
            }
        },
        {
            "name": "list_activity",
            "description": "Get recent activity for a project",
            "parameters": {
                "type": "object",
                "properties": {
                    "projectId": {"type": "string", "description": "Project ID"},
                    "limit": {"type": "number", "default": 20}
                },
                "required": ["projectId"]
            }
        },
        {
            "name": "memory_search",
            "description": "Semantic search across project memory",
            "parameters": {
                "type": "object",
                "properties": {
                    "projectId": {"type": "string", "description": "Project ID"},
                    "query": {"type": "string", "description": "Search query"}
                },
                "required": ["projectId", "query"]
            }
        }
    ]

    # Build system message with project context
    system_message = f"""
    You are an AI assistant helping users understand their project memory.
    Current Project ID: {project_id}

    You have access to these tools:
    - list_projects: List user's projects
    - list_tasks: List tasks in a project
    - list_activity: Get recent activity
    - memory_search: Search project memory semantically

    Use these tools to answer user questions accurately.
    """

    # Convert messages to Gemini format
    chat_messages = []
    for msg in messages:
        if msg["role"] == "system":
            system_message = msg["content"]
        else:
            chat_messages.append(msg)

    # Build the prompt with tool descriptions
    tool_prompt = f"""
{system_message}

To use tools, format your response as:
TOOL: tool_name
ARGS: {{"arg1": "value1"}}

Available tools:
{json.dumps(tools, indent=2)}
"""

    # Add system context to first message
    full_messages = [{"role": "user", "content": tool_prompt}] + chat_messages

    # Convert to Gemini chat format
    chat_history = []
    for msg in full_messages[:-1]:  # All except last
        chat_history.append({
            "role": "user" if msg["role"] == "user" else "model",
            "parts": [msg["content"]]
        })

    # Get best available model from router for chat
    model_name = router.get_model_for_task("chat")
    if not model_name:
        raise Exception("All models are rate limited. Please try again in a minute.")

    model = router.models[model_name]
    router._record_usage(model_name)

    # Start chat session with selected model
    chat = model.start_chat(history=chat_history)

    # Send last message
    last_message = full_messages[-1]["content"]
    response = chat.send_message(last_message)

    # Check if response contains tool call
    response_text = response.text

    # Simple tool detection
    if "TOOL:" in response_text and "ARGS:" in response_text:
        # Parse tool call
        lines = response_text.split("\n")
        tool_name = None
        args = None

        for line in lines:
            if line.startswith("TOOL:"):
                tool_name = line.replace("TOOL:", "").strip()
            elif line.startswith("ARGS:"):
                args = json.loads(line.replace("ARGS:", "").strip())

        # Execute tool if found
        if tool_name and args:
            tool_result = None

            if tool_name == "list_projects":
                tool_result = list_projects(user_id=args["userId"])
            elif tool_name == "list_tasks":
                tool_result = list_tasks(
                    project_id=args["projectId"],
                    status=args.get("status")
                )
            elif tool_name == "list_activity":
                tool_result = list_activity(
                    project_id=args["projectId"],
                    limit=args.get("limit", 20)
                )
            elif tool_name == "memory_search":
                tool_result = await memory_search(
                    project_id=args["projectId"],
                    query=args["query"]
                )

            # Send tool result back to model
            if tool_result:
                follow_up = f"Tool {tool_name} returned: {json.dumps(tool_result)}\n\nBased on this, answer the user's question."
                final_response = chat.send_message(follow_up)
                return final_response.text

    return response_text


async def task_chat(
    task_id: str,
    task_title: str,
    task_description: str,
    project_id: str,
    user_id: str,
    message: str,
    history: list[dict],
    current_datetime: str
) -> dict:
    """Chat with AI agent while working on a task.

    The agent can:
    - Answer questions and give coding advice
    - Search project memory for context
    - Complete the task when user indicates they're done

    Args:
        task_id: ID of the task being worked on
        task_title: Title of the task
        task_description: Description of the task
        project_id: Project ID
        user_id: User ID working on the task
        message: User's message
        history: Conversation history
        current_datetime: Current timestamp

    Returns:
        {message: str, taskCompleted?: bool, taskStatus?: str}
    """
    from app.tools.memory import complete_task, memory_search
    from app.model_router import router

    # System prompt with task context
    system_prompt = f"""You are an AI assistant helping a developer work on a task.

CURRENT TASK:
- Title: {task_title}
- Description: {task_description or 'No description'}
- Task ID: {task_id}

USER: {user_id}
PROJECT: {project_id}
CURRENT TIME: {current_datetime}

YOUR CAPABILITIES:
1. Answer questions and give coding advice related to the task
2. Search project memory for relevant context (completed tasks, documentation)
3. Complete the task when the user EXPLICITLY CONFIRMS

TASK COMPLETION FLOW:
When the user indicates they've finished (e.g., "I'm done", "finished it", describes what they did):
1. Briefly acknowledge what they accomplished
2. ASK SIMPLY: "Would you like me to mark this task as complete?" (just this question, nothing more)
3. WAIT for user confirmation (e.g., "yes", "mark it", "complete it", "sure")
4. ONLY after explicit confirmation, call the complete_task tool

IMPORTANT:
- Do NOT call complete_task until the user explicitly confirms
- Do NOT ask for additional details or descriptions when confirming - just ask yes/no
- The user has already told you what they did - use that information for the complete_task tool

To use tools, format your response as:
TOOL: tool_name
ARGS: {{"arg1": "value1"}}
RESULT_PENDING

After I provide the tool result, give your final response to the user.

Available tools:
- memory_search: Search project memory. Args: {{"query": "search terms"}}
- complete_task: Mark task as complete. Args: {{"what_i_did": "description of work done", "code_snippet": "optional code"}}

Be helpful, concise, and focused on helping complete the task."""

    # Build conversation for the model
    chat_messages = []

    # Add history (convert role names)
    for msg in history:
        role = "model" if msg["role"] == "assistant" else "user"
        chat_messages.append({
            "role": role,
            "parts": [msg["content"]]
        })

    # Get best available model
    model_name = router.get_model_for_task("chat")
    if not model_name:
        return {"message": "All AI models are temporarily unavailable. Please try again in a minute."}

    model = router.models[model_name]
    router._record_usage(model_name)

    # Start chat with system context in first message
    first_message = f"{system_prompt}\n\nUser's first message will follow."
    chat_history = [{"role": "user", "parts": [first_message]}, {"role": "model", "parts": ["Understood. I'm ready to help you work on this task. What would you like to know or do?"]}]

    # Add conversation history
    chat_history.extend(chat_messages)

    chat = model.start_chat(history=chat_history)

    # Send user's message
    response = chat.send_message(message)
    response_text = response.text

    # Check for tool calls
    task_completed = False
    task_status = "in_progress"

    if "TOOL:" in response_text and "ARGS:" in response_text:
        lines = response_text.split("\n")
        tool_name = None
        args = None

        for line in lines:
            if line.startswith("TOOL:"):
                tool_name = line.replace("TOOL:", "").strip()
            elif line.startswith("ARGS:"):
                try:
                    args = json.loads(line.replace("ARGS:", "").strip())
                except json.JSONDecodeError:
                    continue

        if tool_name and args:
            tool_result = None

            if tool_name == "memory_search":
                tool_result = await memory_search(
                    project_id=project_id,
                    query=args.get("query", "")
                )
            elif tool_name == "complete_task":
                what_i_did = args.get("what_i_did", message)
                code_snippet = args.get("code_snippet")

                tool_result = await complete_task(
                    task_id=task_id,
                    project_id=project_id,
                    user_id=user_id,
                    what_i_did=what_i_did,
                    code_snippet=code_snippet
                )

                if "error" not in tool_result:
                    task_completed = True
                    task_status = "done"

            # Get follow-up response with tool result
            if tool_result:
                follow_up = f"Tool {tool_name} returned: {json.dumps(tool_result)}\n\nProvide your response to the user."
                final_response = chat.send_message(follow_up)
                response_text = final_response.text

    return {
        "message": response_text,
        "taskCompleted": task_completed,
        "taskStatus": task_status
    }