"""
Embedding router for managing vector embeddings
"""

import asyncio
from fastapi import APIRouter, HTTPException, Depends
from typing import Dict, Any

from models.embedding import (
    EmbedRequest, 
    EmbedResponse, 
    EmbedAllRequest, 
    EmbedAllResponse
)
from core.dependencies import get_vector_db, get_task_manager, get_logger

router = APIRouter(prefix="/embedding", tags=["embedding"])


@router.post("/markdown", response_model=EmbedResponse)
async def embed_markdown(
    request: EmbedRequest,
    vector_db=Depends(get_vector_db),
    task_manager=Depends(get_task_manager),
    logger=Depends(get_logger)
):
    """Embed markdown files from a specific directory into ChromaDB"""
    try:
        # Run the embedding process in a background task
        task_id = await task_manager.create_task(
            asyncio.to_thread(
                vector_db.embed_markdown_files,
                request.directory,
                request.collection_name,
                request.chunk_size,
                request.chunk_overlap
            )
        )
        
        # Wait for the task to complete
        task_status = await task_manager.wait_for_task(task_id, timeout=60.0)
        
        if task_status["status"] == "completed":
            documents_embedded = task_status.get("result", 0)
            return EmbedResponse(
                status="success",
                message=f"Successfully embedded {documents_embedded} documents",
                documents_embedded=documents_embedded
            )
        elif task_status["status"] == "timeout":
            return EmbedResponse(
                status="pending",
                message="Embedding process started but still running. Check back later.",
                documents_embedded=0
            )
        else:
            return EmbedResponse(
                status="error",
                message=f"Error embedding documents: {task_status.get('error', 'Unknown error')}",
                documents_embedded=0
            )
    except Exception as e:
        logger.error(f"Error in embed_markdown: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/all", response_model=EmbedAllResponse)
async def embed_all_markdown(
    request: EmbedAllRequest,
    vector_db=Depends(get_vector_db),
    task_manager=Depends(get_task_manager),
    logger=Depends(get_logger)
):
    """Embed all markdown files (pages and posts) into ChromaDB"""
    try:
        # Define directories
        pages_dir = "markdown-data/pages_markdown"
        posts_dir = "markdown-data/posts_markdown"
        
        # Process pages markdown files
        pages_task_id = await task_manager.create_task(
            asyncio.to_thread(
                vector_db.embed_markdown_files,
                pages_dir,
                request.collection_name,
                request.chunk_size,
                request.chunk_overlap
            )
        )
        
        # Wait for pages task to complete
        pages_status = await task_manager.wait_for_task(pages_task_id, timeout=120.0)
        pages_embedded = pages_status.get("result", 0) if pages_status["status"] == "completed" else 0
        
        # Process posts markdown files
        posts_task_id = await task_manager.create_task(
            asyncio.to_thread(
                vector_db.embed_markdown_files,
                posts_dir,
                request.collection_name,
                request.chunk_size,
                request.chunk_overlap
            )
        )
        
        # Wait for posts task to complete
        posts_status = await task_manager.wait_for_task(posts_task_id, timeout=120.0)
        posts_embedded = posts_status.get("result", 0) if posts_status["status"] == "completed" else 0
        
        # Calculate total embedded documents
        total_embedded = pages_embedded + posts_embedded
        
        # Determine overall status
        if pages_status["status"] == "completed" and posts_status["status"] == "completed":
            status = "success"
            message = f"Successfully embedded all markdown files into collection {request.collection_name}"
        elif pages_status["status"] == "timeout" or posts_status["status"] == "timeout":
            status = "partial"
            message = "Embedding process started but some tasks are still running. Check back later."
        else:
            status = "error"
            message = "Error embedding some markdown files"
        
        # Return the results
        return EmbedAllResponse(
            status=status,
            message=message,
            pages_embedded=pages_embedded,
            posts_embedded=posts_embedded,
            total_embedded=total_embedded,
            embedding_model=request.embed_backend,
            collection_name=request.collection_name
        )
    except Exception as e:
        logger.error(f"Error in embed_all_markdown: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))