from fastapi import APIRouter, HTTPException
from app.core.logger import logger
from app.models.chunk import (
    EmbeddingJobRequest,
    EmbeddingJobResponse,
    BatchEmbeddingRequest,
    BatchEmbeddingResponse,
)
from app.models.events import EmbeddingJobDispatchEvent
from app.services.embedding_dispatch_service import to_embedding_task_request
from app.workers.celery_app import celery_app
from app.workers.embedding_tasks import (
    embed_chunk_task,
    embed_chunks_batch_task,
)
from app.core.config import settings

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


@router.post("/embed", response_model=EmbeddingJobResponse)
async def queue_embedding(request: EmbeddingJobRequest) -> EmbeddingJobResponse:
    """
    Queue a chunk for embedding.
    
    - **chunk_id**: Unique identifier for the chunk
    - **chunk_text**: Text content to embed
    - **source_type**: Type of source (code, repo_doc, commit, etc.)
    """
    try:
        logger.info(f"Queueing embedding for chunk {request.chunk_id}")
        
        # Queue task in Celery
        task = embed_chunk_task.delay(
            chunk_id=request.chunk_id,
            chunk_text=request.chunk_text,
            metadata=request.metadata,
        )
        
        return EmbeddingJobResponse(
            task_id=task.id,
            chunk_id=request.chunk_id,
            status="queued",
            message="Embedding job queued successfully",
        )
    except Exception as e:
        logger.error(f"Failed to queue embedding: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/dispatch")
async def queue_embedding_dispatch(event: EmbeddingJobDispatchEvent) -> dict:
    """
    Receive an embedding dispatch event from docs-ingestion and queue one chunk.
    """
    try:
        task_request = to_embedding_task_request(event)
        logger.info(f"Queueing embedding dispatch job {event.job_id} for chunk {task_request['chunk_id']}")

        task = embed_chunk_task.delay(
            chunk_id=task_request["chunk_id"],
            chunk_text=task_request["chunk_text"],
            metadata=task_request["metadata"],
        )

        return {
            "job_id": task_request["job_id"],
            "task_id": task.id,
            "chunk_id": task_request["chunk_id"],
            "status": "queued",
        }
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to queue embedding dispatch: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/embed-batch", response_model=BatchEmbeddingResponse)
async def queue_batch_embedding(request: BatchEmbeddingRequest) -> BatchEmbeddingResponse:
    """
    Queue multiple chunks for embedding in batch.
    
    - **chunks**: List of chunk objects to embed
    - **source_type**: Type of source for all chunks
    """
    try:
        if not request.chunks or len(request.chunks) == 0:
            raise ValueError("No chunks provided")
        
        logger.info(f"Queueing batch embedding for {len(request.chunks)} chunks")
        
        chunk_ids = [chunk.chunk_id for chunk in request.chunks]
        chunk_texts = [chunk.chunk_text for chunk in request.chunks]
        
        # Queue batch task
        task = embed_chunks_batch_task.delay(
            chunk_ids=chunk_ids,
            chunk_texts=chunk_texts,
            source_type=request.source_type,
        )
        
        return BatchEmbeddingResponse(
            batch_id=task.id,
            total_chunks=len(request.chunks),
            queued_chunks=len(request.chunks),
            status="queued",
        )
    except Exception as e:
        logger.error(f"Failed to queue batch embedding: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/embed/{task_id}")
async def get_embedding_status(task_id: str) -> dict:
    """
    Get status of an embedding job.
    
    - **task_id**: Task ID returned from queue endpoint
    """
    try:
        task_result = celery_app.AsyncResult(task_id)
        
        if task_result.state == "PENDING":
            return {
                "task_id": task_id,
                "status": "pending",
                "result": None,
            }
        elif task_result.state == "SUCCESS":
            return {
                "task_id": task_id,
                "status": "completed",
                "result": task_result.result,
            }
        elif task_result.state == "FAILURE":
            return {
                "task_id": task_id,
                "status": "failed",
                "error": str(task_result.info),
            }
        else:  # RETRY or other states
            return {
                "task_id": task_id,
                "status": task_result.state.lower(),
                "result": task_result.result if task_result.successful() else None,
            }
    except Exception as e:
        logger.error(f"Failed to get task status: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/embed-sync")
async def embed_text_synchronous(request: EmbeddingJobRequest) -> dict:
    """
    Embed text synchronously (no queueing).
    Warning: This blocks the request, use for small texts only.
    
    - **chunk_text**: Text to embed
    """
    try:
        logger.info(f"Embedding text synchronously for chunk {request.chunk_id}")
        
        from app.services.embedding_service import get_embedding_service
        embedding_service = get_embedding_service()
        embedding = await embedding_service.embed_text(request.chunk_text)
        
        return {
            "chunk_id": request.chunk_id,
            "embedding": embedding,
            "status": "completed",
        }
    except Exception as e:
        logger.error(f"Synchronous embedding failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/health")
async def embedding_service_health() -> dict:
    """Lightweight health check without calling external providers."""
    return {
        "status": "healthy",
        "service": "embedding",
        "provider": settings.EMBEDDING_PROVIDER,
        "model": settings.EMBEDDING_MODEL,
        "queue": settings.CELERY_BROKER_URL,
    }