foxus/backend/app/api/ai.py

"""
AI API routes for code assistance
"""

import time
from typing import List
from fastapi import APIRouter, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse
from app.models.ai import (
    AIRequest, AIResponse, CodeCompletionRequest, CodeCompletionResponse,
    ExplainRequest, RefactorRequest, FixRequest, MultiFileRequest
)
from app.services.ollama_service import ollama_service

router = APIRouter()

@router.post("/process", response_model=AIResponse)
async def process_ai_request(request: AIRequest):
    """Process general AI request for code assistance"""
    start_time = time.time()
    
    try:
        # Check if Ollama is available
        if not await ollama_service.is_available():
            raise HTTPException(
                status_code=503,
                detail="Ollama service is not available. Please ensure Ollama is running."
            )
        
        # Build prompt based on command
        prompt = ollama_service.build_prompt(
            command=request.command,
            code=request.code,
            language=request.language,
            context=request.context
        )
        
        # Generate completion
        result = await ollama_service.generate_completion(
            prompt=prompt,
            model=request.model
        )
        
        execution_time = time.time() - start_time
        
        return AIResponse(
            success=True,
            result=result,
            execution_time=execution_time,
            model_used=request.model or ollama_service.default_model
        )
        
    except Exception as e:
        execution_time = time.time() - start_time
        return AIResponse(
            success=False,
            result=f"Error processing request: {str(e)}",
            execution_time=execution_time
        )

@router.post("/explain", response_model=AIResponse)
async def explain_code(request: ExplainRequest):
    """Explain code functionality"""
    start_time = time.time()
    
    try:
        if not await ollama_service.is_available():
            raise HTTPException(status_code=503, detail="Ollama service unavailable")
        
        # Build explanation prompt
        prompt = f"""
Explain the following {request.language.value if request.language else 'code'} code in {request.detail_level} detail:

```{request.language.value if request.language else 'code'}
{request.code}
```

Please provide a clear explanation that covers:
1. What this code does
2. How it works
3. Key concepts used
4. Any potential issues or improvements

Explanation:"""
        
        result = await ollama_service.generate_completion(prompt=prompt)
        execution_time = time.time() - start_time
        
        return AIResponse(
            success=True,
            result=result,
            execution_time=execution_time,
            model_used=ollama_service.default_model
        )
        
    except Exception as e:
        execution_time = time.time() - start_time
        return AIResponse(
            success=False,
            result=f"Error explaining code: {str(e)}",
            execution_time=execution_time
        )

@router.post("/refactor", response_model=AIResponse)
async def refactor_code(request: RefactorRequest):
    """Refactor code for better quality"""
    start_time = time.time()
    
    try:
        if not await ollama_service.is_available():
            raise HTTPException(status_code=503, detail="Ollama service unavailable")
        
        prompt = f"""
Refactor the following {request.language.value if request.language else 'code'} code to improve readability, performance, and maintainability:

```{request.language.value if request.language else 'code'}
{request.code}
```

Focus on {request.refactor_type} improvements. Please provide:
1. The refactored code
2. Explanation of changes made
3. Benefits of the refactoring

Refactored code:"""
        
        result = await ollama_service.generate_completion(prompt=prompt)
        execution_time = time.time() - start_time
        
        return AIResponse(
            success=True,
            result=result,
            execution_time=execution_time,
            model_used=ollama_service.default_model
        )
        
    except Exception as e:
        execution_time = time.time() - start_time
        return AIResponse(
            success=False,
            result=f"Error refactoring code: {str(e)}",
            execution_time=execution_time
        )

@router.post("/fix", response_model=AIResponse)
async def fix_code(request: FixRequest):
    """Fix bugs in code"""
    start_time = time.time()
    
    try:
        if not await ollama_service.is_available():
            raise HTTPException(status_code=503, detail="Ollama service unavailable")
        
        prompt = ollama_service.build_prompt(
            command="fix",
            code=request.code,
            language=request.language,
            error_message=request.error_message
        )
        
        result = await ollama_service.generate_completion(prompt=prompt)
        execution_time = time.time() - start_time
        
        return AIResponse(
            success=True,
            result=result,
            execution_time=execution_time,
            model_used=ollama_service.default_model
        )
        
    except Exception as e:
        execution_time = time.time() - start_time
        return AIResponse(
            success=False,
            result=f"Error fixing code: {str(e)}",
            execution_time=execution_time
        )

@router.post("/complete", response_model=CodeCompletionResponse)
async def complete_code(request: CodeCompletionRequest):
    """Generate code completion"""
    start_time = time.time()
    
    try:
        if not await ollama_service.is_available():
            raise HTTPException(status_code=503, detail="Ollama service unavailable")
        
        # Extract context around cursor position
        code_before = request.code[:request.cursor_position]
        code_after = request.code[request.cursor_position:]
        
        prompt = f"""
Complete the following {request.language.value if request.language else 'code'} code at the cursor position:

```{request.language.value if request.language else 'code'}
{code_before}<CURSOR>{code_after}
```

Provide only the code that should be inserted at the cursor position. Keep it concise and contextually appropriate.

Completion:"""
        
        result = await ollama_service.generate_completion(
            prompt=prompt,
            max_tokens=request.max_tokens
        )
        
        execution_time = time.time() - start_time
        
        return CodeCompletionResponse(
            success=True,
            completions=[result.strip()],
            cursor_position=request.cursor_position,
            execution_time=execution_time
        )
        
    except Exception as e:
        execution_time = time.time() - start_time
        return CodeCompletionResponse(
            success=False,
            completions=[],
            cursor_position=request.cursor_position,
            execution_time=execution_time
        )

@router.post("/stream")
async def stream_ai_response(request: AIRequest):
    """Stream AI response for real-time feedback"""
    try:
        if not await ollama_service.is_available():
            raise HTTPException(status_code=503, detail="Ollama service unavailable")
        
        prompt = ollama_service.build_prompt(
            command=request.command,
            code=request.code,
            language=request.language,
            context=request.context
        )
        
        async def generate():
            try:
                async for chunk in ollama_service.generate_streaming(
                    prompt=prompt,
                    model=request.model
                ):
                    yield f"data: {chunk}\n\n"
                yield "data: [DONE]\n\n"
            except Exception as e:
                yield f"data: ERROR: {str(e)}\n\n"
        
        return StreamingResponse(
            generate(),
            media_type="text/plain",
            headers={"Cache-Control": "no-cache"}
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/multifile", response_model=AIResponse)
async def process_multifile_request(request: MultiFileRequest):
    """Process multi-file analysis request"""
    start_time = time.time()
    
    try:
        if not await ollama_service.is_available():
            raise HTTPException(status_code=503, detail="Ollama service unavailable")
        
        # Build context from multiple files
        file_context = "\n\n".join([
            f"File: {file_info['path']}\n```\n{file_info['content']}\n```"
            for file_info in request.files
        ])
        
        # Focus on specific file if provided
        focus_context = ""
        if request.focus_file:
            focus_file = next(
                (f for f in request.files if f['path'] == request.focus_file),
                None
            )
            if focus_file:
                focus_context = f"\n\nFocus on this file: {focus_file['path']}\n```\n{focus_file['content']}\n```"
        
        prompt = f"""
Analyze the following multi-file codebase and {request.command.value}:

{file_context}
{focus_context}

{f"Additional context: {request.context}" if request.context else ""}

Please provide analysis considering the relationships between files and overall code structure.

Response:"""
        
        result = await ollama_service.generate_completion(prompt=prompt)
        execution_time = time.time() - start_time
        
        return AIResponse(
            success=True,
            result=result,
            execution_time=execution_time,
            model_used=ollama_service.default_model,
            metadata={"files_analyzed": len(request.files)}
        )
        
    except Exception as e:
        execution_time = time.time() - start_time
        return AIResponse(
            success=False,
            result=f"Error processing multifile request: {str(e)}",
            execution_time=execution_time
        )

@router.get("/health")
async def ai_health():
    """Check AI service health"""
    is_available = await ollama_service.is_available()
    models = await ollama_service.list_models() if is_available else []
    
    return {
        "ollama_available": is_available,
        "models_count": len(models),
        "default_model": ollama_service.default_model,
        "base_url": ollama_service.base_url
    }