backend/app/main.py

"""FastAPI application for Project Memory - API layer calling MCP tools."""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from contextlib import asynccontextmanager
from typing import List, Optional
import os
import asyncio
from dotenv import load_dotenv

from app import schemas
from app.schemas import (
    ProjectCreate, ProjectJoin, Project,
    TaskCreate, Task, TaskCompleteRequest, TaskCompleteResponse,
    ActivityResponse, SearchRequest, SearchResponse,
    SmartQueryRequest, SmartQueryResponse,
    ChatRequest, ChatResponse, ErrorResponse,
    UserCreate, User
)

from app.tools.projects import create_project, list_projects, join_project, check_project_id_available
from app.tools.tasks import create_task, list_tasks, list_activity
from app.tools.memory import complete_task, memory_search

# Load environment variables
load_dotenv()


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize database and vector store on startup."""
    from app.database import init_db, SessionLocal
    from app.vectorstore import init_vectorstore
    from app.models import User, AI_AGENT_USER_ID
    from app.agent_worker import agent_loop

    init_db()
    init_vectorstore()
    print("[OK] Database and vector store initialized")

    # Ensure AI Agent user exists
    db = SessionLocal()
    try:
        if not db.query(User).filter(User.id == AI_AGENT_USER_ID).first():
            db.add(User(id=AI_AGENT_USER_ID, first_name="AI", last_name="Agent"))
            db.commit()
            print("[OK] AI Agent user created")
    finally:
        db.close()

    # Start agent worker in background
    agent_task = asyncio.create_task(agent_loop())
    print("[OK] Agent worker started")

    yield

    # Cleanup on shutdown
    agent_task.cancel()
    try:
        await agent_task
    except asyncio.CancelledError:
        pass
    print("[OK] Agent worker stopped")


# Initialize FastAPI app
app = FastAPI(
    title="Project Memory API",
    description="Multi-user, multi-project AI memory system powered by MCP",
    version="1.0.0",
    lifespan=lifespan
)

# Configure CORS
frontend_url = os.getenv("FRONTEND_URL", "http://localhost:5173")
app.add_middleware(
    CORSMiddleware,
    allow_origins=[frontend_url, "http://localhost:5173", "http://localhost:3000"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ==================== Health Check ====================
@app.get("/")
async def root():
    """Health check endpoint."""
    return {"status": "ok", "message": "Project Memory API is running"}


# ==================== User Endpoints ====================
@app.post("/api/users", response_model=schemas.User)
async def create_user(user: schemas.UserCreate):
    """Create a new user."""
    from app.database import get_db
    from app.models import User, generate_user_id
    
    db = next(get_db())
    
    # Generate unique user ID (first 3 letters of firstname + 4 random digits)
    user_id = generate_user_id(user.firstName)
    
    # Ensure ID is unique (regenerate if collision)
    while db.query(User).filter(User.id == user_id).first():
        user_id = generate_user_id(user.firstName)
    
    # Create new user
    new_user = User(
        id=user_id,
        first_name=user.firstName,
        last_name=user.lastName,
        avatar_url=user.avatar_url
    )
    db.add(new_user)
    db.commit()
    db.refresh(new_user)
    
    return {
        "id": new_user.id,
        "firstName": new_user.first_name,
        "lastName": new_user.last_name,
        "avatar_url": new_user.avatar_url,
        "created_at": new_user.created_at
    }


@app.get("/api/users/{user_id}", response_model=schemas.User)
async def get_user(user_id: str):
    """Get user by ID."""
    from app.database import get_db
    from app.models import User
    
    db = next(get_db())
    user = db.query(User).filter(User.id == user_id).first()
    
    if not user:
        raise HTTPException(status_code=404, detail="User not found")
    
    return {
        "id": user.id,
        "firstName": user.first_name,
        "lastName": user.last_name,
        "avatar_url": user.avatar_url,
        "created_at": user.created_at
    }


@app.get("/api/users")
async def list_users():
    """List all users."""
    from app.database import get_db
    from app.models import User
    
    db = next(get_db())
    users = db.query(User).all()
    return [{"id": u.id, "firstName": u.first_name, "lastName": u.last_name} for u in users]


# ==================== Project Endpoints ====================
@app.get("/api/projects/check/{project_id}")
async def check_project_availability(project_id: str):
    """Check if a project ID is available."""
    try:
        result = check_project_id_available(project_id=project_id)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/projects")
async def get_projects(userId: str):
    """List all projects for a user."""
    try:
        result = list_projects(user_id=userId)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/projects", response_model=Project)
async def create_new_project(project: ProjectCreate):
    """Create a new project."""
    try:
        result = create_project(
            name=project.name,
            description=project.description,
            user_id=project.userId
        )
        if "error" in result:
            raise HTTPException(status_code=400, detail=result["error"])
        return result
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/projects/{project_id}/join")
async def join_existing_project(project_id: str, request: ProjectJoin):
    """Join an existing project."""
    try:
        result = join_project(
            project_id=project_id,
            user_id=request.userId
        )
        if "error" in result:
            raise HTTPException(status_code=400, detail=result["error"])
        return result
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/projects/{project_id}/members")
async def get_project_members(project_id: str):
    """Get all members of a project."""
    from app.database import get_db
    from app.models import Project, ProjectMembership, User

    db = next(get_db())
    try:
        # Check project exists
        project = db.query(Project).filter(Project.id == project_id).first()
        if not project:
            raise HTTPException(status_code=404, detail="Project not found")

        # Get all memberships with user details
        memberships = db.query(ProjectMembership, User).join(
            User, ProjectMembership.user_id == User.id
        ).filter(ProjectMembership.project_id == project_id).all()

        members = [
            {
                "id": user.id,
                "firstName": user.first_name,
                "lastName": user.last_name,
                "avatar_url": user.avatar_url,
                "role": membership.role,
                "joined_at": membership.joined_at.isoformat() if membership.joined_at else None
            }
            for membership, user in memberships
        ]

        return {"members": members}
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db.close()


# ==================== Agent Endpoints ====================
@app.post("/api/projects/{project_id}/agent/enable")
async def enable_agent(project_id: str):
    """Enable AI agent for this project. Adds agent to team."""
    from app.database import get_db
    from app.models import Project, ProjectMembership, AI_AGENT_USER_ID

    db = next(get_db())
    try:
        # Check project exists
        project = db.query(Project).filter(Project.id == project_id).first()
        if not project:
            raise HTTPException(status_code=404, detail="Project not found")

        # Enable agent
        project.agent_enabled = True

        # Add agent to project membership if not already
        existing = db.query(ProjectMembership).filter(
            ProjectMembership.project_id == project_id,
            ProjectMembership.user_id == AI_AGENT_USER_ID
        ).first()

        if not existing:
            membership = ProjectMembership(
                project_id=project_id,
                user_id=AI_AGENT_USER_ID,
                role="agent"
            )
            db.add(membership)

        db.commit()
        return {"message": "AI Agent enabled", "project_id": project_id}
    except HTTPException:
        raise
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db.close()


@app.post("/api/projects/{project_id}/agent/disable")
async def disable_agent(project_id: str):
    """Disable AI agent for this project. Removes from team."""
    from app.database import get_db
    from app.models import Project, ProjectMembership, AI_AGENT_USER_ID

    db = next(get_db())
    try:
        # Check project exists
        project = db.query(Project).filter(Project.id == project_id).first()
        if not project:
            raise HTTPException(status_code=404, detail="Project not found")

        # Disable agent
        project.agent_enabled = False

        # Remove agent from project membership
        db.query(ProjectMembership).filter(
            ProjectMembership.project_id == project_id,
            ProjectMembership.user_id == AI_AGENT_USER_ID
        ).delete()

        db.commit()
        return {"message": "AI Agent disabled", "project_id": project_id}
    except HTTPException:
        raise
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db.close()


# ==================== Task Endpoints ====================
@app.get("/api/projects/{project_id}/tasks")
async def get_project_tasks(project_id: str, status: Optional[str] = None):
    """Get all tasks for a project, optionally filtered by status."""
    try:
        result = list_tasks(
            project_id=project_id,
            status=status
        )
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/projects/{project_id}/tasks", response_model=Task)
async def create_new_task(project_id: str, task: TaskCreate):
    """Create a new task in a project."""
    try:
        result = create_task(
            project_id=project_id,
            title=task.title,
            description=task.description,
            assigned_to=task.assignedTo
        )
        if "error" in result:
            raise HTTPException(status_code=400, detail=result["error"])
        return result
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/projects/{project_id}/tasks/generate")
async def generate_project_tasks(project_id: str, request: dict = None):
    """Generate demo tasks for a project using AI.

    Does NOT save to database - returns generated tasks for user to edit.
    Max 50 tasks.
    """
    from app.llm import generate_tasks
    from app.database import get_db
    from app.models import Project

    db = next(get_db())
    try:
        # Get project details
        project = db.query(Project).filter(Project.id == project_id).first()
        if not project:
            raise HTTPException(status_code=404, detail="Project not found")

        # Get count from request, default 50, max 50
        count = min(request.get("count", 50) if request else 50, 50)

        # Generate tasks using LLM (no user prompt needed)
        tasks = await generate_tasks(
            project_name=project.name,
            project_description=project.description,
            count=count
        )

        return {"tasks": tasks}
    except HTTPException:
        raise
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db.close()


@app.post("/api/tasks/{task_id}/complete", response_model=TaskCompleteResponse)
async def complete_existing_task(task_id: str, request: TaskCompleteRequest):
    """Complete a task with documentation. Generates AI docs and stores embeddings."""
    from app.database import get_db
    from app.models import Task as TaskModel

    db = next(get_db())
    try:
        task = db.query(TaskModel).filter(TaskModel.id == task_id).first()

        if not task:
            raise HTTPException(status_code=404, detail="Task not found")

        result = await complete_task(
            task_id=task_id,
            project_id=task.project_id,
            user_id=request.userId,
            what_i_did=request.whatIDid,
            code_snippet=request.codeSnippet
        )

        if "error" in result:
            raise HTTPException(status_code=400, detail=result["error"])

        return result
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db.close()


@app.patch("/api/tasks/{task_id}/status")
async def update_task_status(task_id: str, request: dict):
    """Update task status (for kanban board)."""
    from app.database import get_db
    from app.models import Task as TaskModel, TaskStatus

    db = next(get_db())
    try:
        task = db.query(TaskModel).filter(TaskModel.id == task_id).first()

        if not task:
            raise HTTPException(status_code=404, detail="Task not found")

        new_status = request.get("status")
        user_id = request.get("userId")  # Who is making this change

        if new_status not in ["todo", "in_progress", "done"]:
            raise HTTPException(status_code=400, detail="Invalid status. Must be: todo, in_progress, or done")

        # Update working_by based on status
        if new_status == "in_progress" and user_id:
            task.working_by = user_id
        elif new_status in ["todo", "done"]:
            task.working_by = None  # Clear when not in progress

        task.status = TaskStatus(new_status)
        db.commit()
        db.refresh(task)

        return {
            "id": str(task.id),
            "project_id": task.project_id,
            "title": task.title,
            "description": task.description,
            "status": task.status.value,
            "assigned_to": task.assigned_to,
            "working_by": task.working_by,
            "created_at": task.created_at.isoformat() if task.created_at else None
        }
    except HTTPException:
        raise
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db.close()


@app.post("/api/tasks/{task_id}/chat")
async def chat_with_task_agent(task_id: str, request: dict):
    """Chat with AI agent while working on a task.

    The agent can answer questions, search project memory, and complete tasks.
    """
    from app.database import get_db
    from app.models import Task as TaskModel
    from app.llm import task_chat

    db = next(get_db())
    try:
        # Get task details
        task = db.query(TaskModel).filter(TaskModel.id == task_id).first()
        if not task:
            raise HTTPException(status_code=404, detail="Task not found")

        # Extract request data
        project_id = request.get("projectId", task.project_id)
        user_id = request.get("userId")
        message = request.get("message")
        history = request.get("history", [])
        current_datetime = request.get("currentDatetime", "")

        if not user_id:
            raise HTTPException(status_code=400, detail="userId is required")
        if not message:
            raise HTTPException(status_code=400, detail="message is required")

        # Call the task chat function
        result = await task_chat(
            task_id=task_id,
            task_title=task.title,
            task_description=task.description or "",
            project_id=project_id,
            user_id=user_id,
            message=message,
            history=history,
            current_datetime=current_datetime
        )

        return result
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db.close()


# ==================== Activity Feed Endpoint ====================
@app.get("/api/projects/{project_id}/activity")
async def get_project_activity(project_id: str, limit: int = 20):
    """Get recent activity for a project."""
    try:
        result = list_activity(
            project_id=project_id,
            limit=limit
        )
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ==================== Search Endpoint ====================
@app.post("/api/projects/{project_id}/search", response_model=SearchResponse)
async def search_project_memory(project_id: str, request: SearchRequest):
    """Semantic search across project memory."""
    try:
        result = await memory_search(
            project_id=project_id,
            query=request.query,
            filters=request.filters.dict() if request.filters else None
        )
        
        if "error" in result:
            raise HTTPException(status_code=400, detail=result["error"])
        
        return result
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ==================== Smart Query Endpoint ====================
@app.post("/api/projects/{project_id}/smart-query", response_model=SmartQueryResponse)
async def smart_query_project(project_id: str, request: SmartQueryRequest):
    """Natural language query with context awareness.

    Understands queries like:
    - "What did I do yesterday?"
    - "What did Alice do today?"
    - "How does the auth system work?"
    - "Task 13 status?"
    """
    try:
        from app.smart_query import smart_query

        result = await smart_query(
            project_id=project_id,
            query=request.query,
            current_user_id=request.currentUserId,
            current_datetime=request.currentDatetime
        )

        if "error" in result.get("answer", ""):
            raise HTTPException(status_code=400, detail=result["answer"])

        return result
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ==================== Chat Endpoint ====================
@app.post("/api/chat", response_model=ChatResponse)
async def chat_with_ai(request: ChatRequest):
    """Chat with AI using MCP tools."""
    try:
        # Import here to avoid circular dependency
        from app.llm import chat_with_tools
        
        # Convert messages to dict format
        messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
        
        # Call the chat function with tool support
        result = await chat_with_tools(
            messages=messages,
            project_id=request.projectId
        )
        
        return {"message": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ==================== Error Handlers ====================
@app.exception_handler(HTTPException)
async def http_exception_handler(request, exc):
    """Custom HTTP exception handler."""
    return JSONResponse(
        status_code=exc.status_code,
        content={
            "error": exc.detail,
            "status_code": exc.status_code
        }
    )


@app.exception_handler(Exception)
async def general_exception_handler(request, exc):
    """General exception handler."""
    return JSONResponse(
        status_code=500,
        content={
            "error": "Internal server error",
            "detail": str(exc)
        }
    )


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=8000,
        reload=True
    )