"""Retrieval learning route handlers."""
from __future__ import annotations

import asyncio
from dataclasses import asdict
from time import perf_counter
from typing import TYPE_CHECKING, Any

from logos_adk.observe import Metric, get_observability_store

from ...schemas import (
    RetrievalSourceScoresResponse,
    RetrievalSourceDocumentsResponse,
    RetrievalReindexTriggersResponse,
    RetrievalTriggerActionResponse,
    RetrievalForceReindexResponse,
    RetrievalReindexRequestsResponse,
    RetrievalReindexRequestResponse,
    RetrievalOpsSummaryResponse,
    RetrievalProfilesResponse,
)

if TYPE_CHECKING:
    pass


_ACTIVE_REINDEX_TASKS: set[str] = set()


def get_retrieval_engine(
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> Any | None:
    """Get retrieval learning engine from learning store."""
    if learning_store is None:
        learning_engine = None
    elif hasattr(learning_store, "retrieval"):
        learning_engine = learning_store.retrieval
    elif hasattr(learning_store, "retrieval_learning"):
        learning_engine = learning_store.retrieval_learning
    else:
        learning_engine = None

    if learning_engine is not None:
        return learning_engine

    knowledge_base = getattr(getattr(agent, "_config", None), "knowledge_base", None)
    if knowledge_base is None:
        return None
    if hasattr(knowledge_base, "retrieval_learning"):
        return knowledge_base.retrieval_learning

    return None


def _request_to_dict(request: Any) -> dict[str, Any]:
    return {
        "id": request.id,
        "workspace_id": request.workspace_id,
        "source_key": request.source_key,
        "reason": request.reason,
        "requested_by": request.requested_by,
        "status": request.status,
        "metadata": request.metadata,
        "created_at": request.created_at,
        "updated_at": request.updated_at,
    }


def _reindex_metric_attributes(request: Any, agent: Any | None) -> dict[str, Any]:
    return {
        "agent_name": getattr(agent, "name", "unknown"),
        "workspace_id": request.workspace_id or "__global__",
        "source_key": request.source_key,
        "request_id": request.id,
        "status": request.status,
    }


async def _process_reindex_request(
    *,
    engine: Any,
    request_id: str,
    agent: Any | None,
) -> None:
    observability = get_observability_store()
    started_at = perf_counter()
    request = engine.get_reindex_request(request_id) if hasattr(engine, "get_reindex_request") else None
    if request is None or request.status != "queued":
        return

    try:
        if hasattr(engine, "start_reindex_request"):
            request = engine.start_reindex_request(request_id) or request
        documents = engine.list_source_documents(
            workspace_id=request.workspace_id,
            source_key=request.source_key,
        ) if hasattr(engine, "list_source_documents") else []
        if not documents:
            request = engine.fail_reindex_request(
                request_id,
                metadata={"error": "source_snapshot_missing"},
            ) or request
            observability.record_metric(
                Metric(
                    name="adk_reindex_failures_total",
                    value=1.0,
                    unit="count",
                    attributes=_reindex_metric_attributes(request, agent),
                )
            )
        else:
            request = engine.complete_reindex_request(
                request_id,
                metadata={
                    "documents_reindexed": len(documents),
                    "source_document_ids": [document.document_id for document in documents],
                },
            ) or request
        duration_ms = (perf_counter() - started_at) * 1000.0
        observability.record_metric(
            Metric(
                name="adk_reindex_duration_ms",
                value=duration_ms,
                unit="ms",
                attributes=_reindex_metric_attributes(request, agent),
            )
        )
        observability.log(
            level="info" if request.status == "completed" else "error",
            message=f"Retrieval reindex request {request.id} {request.status}",
            source="retrieval.reindex",
            agent_name=getattr(agent, "name", None),
            attributes=_reindex_metric_attributes(request, agent),
        )
    except Exception as exc:
        failed_request = engine.fail_reindex_request(
            request_id,
            metadata={"error": str(exc)},
        ) if hasattr(engine, "fail_reindex_request") else None
        request = failed_request or request
        duration_ms = (perf_counter() - started_at) * 1000.0
        observability.record_metric(
            Metric(
                name="adk_reindex_failures_total",
                value=1.0,
                unit="count",
                attributes=_reindex_metric_attributes(request, agent),
            )
        )
        observability.record_metric(
            Metric(
                name="adk_reindex_duration_ms",
                value=duration_ms,
                unit="ms",
                attributes=_reindex_metric_attributes(request, agent),
            )
        )
        observability.log(
            level="error",
            message=f"Retrieval reindex request {request_id} failed",
            source="retrieval.reindex",
            agent_name=getattr(agent, "name", None),
            attributes={**_reindex_metric_attributes(request, agent), "error": str(exc)},
        )


def _schedule_reindex_request(
    *,
    engine: Any,
    request_id: str,
    agent: Any | None,
) -> None:
    if request_id in _ACTIVE_REINDEX_TASKS:
        return
    _ACTIVE_REINDEX_TASKS.add(request_id)

    async def _runner() -> None:
        try:
            await asyncio.sleep(0)
            await _process_reindex_request(engine=engine, request_id=request_id, agent=agent)
        finally:
            _ACTIVE_REINDEX_TASKS.discard(request_id)

    asyncio.create_task(_runner())


async def get_source_scores_handler(
    workspace_id: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalSourceScoresResponse:
    """Handle GET /learning/retrieval/source-scores."""
    engine = get_retrieval_engine(learning_store, agent)
    
    if engine is None:
        return RetrievalSourceScoresResponse(
            workspace_id=workspace_id,
            items=[],
        )
    
    try:
        if hasattr(engine, "list_source_scores"):
            scores = engine.list_source_scores(workspace_id=workspace_id)
            items = [
                {
                    "workspace_id": s.workspace_id,
                    "source_key": s.source_key,
                    "connector": s.connector,
                    "usefulness_score": s.usefulness_score,
                    "helpful_count": s.helpful_count,
                    "harmful_count": s.harmful_count,
                    "accepted_count": s.accepted_count,
                    "rejected_count": s.rejected_count,
                    "last_updated": s.last_updated,
                }
                for s in scores
            ]
        else:
            items = []
        
        return RetrievalSourceScoresResponse(
            workspace_id=workspace_id,
            items=items,
        )
    except Exception:
        return RetrievalSourceScoresResponse(
            workspace_id=workspace_id,
            items=[],
        )


async def export_source_scores_handler(
    format: str = "json",
    workspace_id: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
):
    """Handle GET /learning/retrieval/source-scores/export."""
    from fastapi.responses import Response
    
    engine = get_retrieval_engine(learning_store, agent)
    
    if engine is None or not hasattr(engine, "export_source_scores_report"):
        return Response(content="[]", media_type="application/json")
    
    try:
        report = engine.export_source_scores_report(
            workspace_id=workspace_id,
            format=format,
        )
        media_type = "text/csv" if format == "csv" else "application/json"
        return Response(content=report, media_type=media_type)
    except Exception:
        return Response(content="[]", media_type="application/json")


async def get_reindex_triggers_handler(
    workspace_id: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalReindexTriggersResponse:
    """Handle GET /learning/retrieval/reindex-triggers."""
    engine = get_retrieval_engine(learning_store, agent)
    
    if engine is None:
        return RetrievalReindexTriggersResponse(
            workspace_id=workspace_id,
            items=[],
        )
    
    try:
        if hasattr(engine, "list_reindex_triggers"):
            triggers = engine.list_reindex_triggers(workspace_id=workspace_id)
            items = [
                {
                    "id": t.id,
                    "workspace_id": t.workspace_id,
                    "source_key": t.source_key,
                    "reason": t.reason,
                    "severity": t.severity,
                    "occurrences": t.occurrences,
                    "status": t.status,
                    "suggested_chunk_size": t.suggested_chunk_size,
                    "suggested_overlap": t.suggested_overlap,
                    "metadata": t.metadata,
                    "created_at": t.created_at,
                    "updated_at": t.updated_at,
                }
                for t in triggers
            ]
        else:
            items = []
        
        return RetrievalReindexTriggersResponse(
            workspace_id=workspace_id,
            items=items,
        )
    except Exception:
        return RetrievalReindexTriggersResponse(
            workspace_id=workspace_id,
            items=[],
        )


async def acknowledge_trigger_handler(
    trigger_id: str,
    note: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalTriggerActionResponse:
    """Handle POST /learning/retrieval/reindex-triggers/{trigger_id}/acknowledge."""
    engine = get_retrieval_engine(learning_store, agent)
    
    if engine is None:
        return RetrievalTriggerActionResponse(trigger={"error": "engine not configured"})
    
    try:
        if hasattr(engine, "acknowledge_reindex_trigger"):
            trigger = engine.acknowledge_reindex_trigger(trigger_id, note=note)
            return RetrievalTriggerActionResponse(
                trigger={
                    "id": trigger.id,
                    "status": trigger.status,
                    "updated_at": trigger.updated_at,
                    "metadata": trigger.metadata,
                } if trigger else {"error": "trigger not found"}
            )
    except Exception:
        pass
    
    return RetrievalTriggerActionResponse(trigger={"status": "acknowledged"})


async def resolve_trigger_handler(
    trigger_id: str,
    note: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalTriggerActionResponse:
    """Handle POST /learning/retrieval/reindex-triggers/{trigger_id}/resolve."""
    engine = get_retrieval_engine(learning_store, agent)
    
    if engine is None:
        return RetrievalTriggerActionResponse(trigger={"error": "engine not configured"})
    
    try:
        if hasattr(engine, "resolve_reindex_trigger"):
            trigger = engine.resolve_reindex_trigger(trigger_id, note=note)
            return RetrievalTriggerActionResponse(
                trigger={
                    "id": trigger.id,
                    "status": trigger.status,
                    "updated_at": trigger.updated_at,
                    "metadata": trigger.metadata,
                } if trigger else {"error": "trigger not found"}
            )
    except Exception:
        pass
    
    return RetrievalTriggerActionResponse(trigger={"status": "resolved"})


async def close_trigger_handler(
    trigger_id: str,
    note: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalTriggerActionResponse:
    """Handle POST /learning/retrieval/reindex-triggers/{trigger_id}/close."""
    engine = get_retrieval_engine(learning_store, agent)
    
    if engine is None:
        return RetrievalTriggerActionResponse(trigger={"error": "engine not configured"})
    
    try:
        if hasattr(engine, "close_reindex_trigger"):
            trigger = engine.close_reindex_trigger(trigger_id, note=note)
            return RetrievalTriggerActionResponse(
                trigger={
                    "id": trigger.id,
                    "status": trigger.status,
                    "updated_at": trigger.updated_at,
                    "metadata": trigger.metadata,
                } if trigger else {"error": "trigger not found"}
            )
    except Exception:
        pass
    
    return RetrievalTriggerActionResponse(trigger={"status": "closed"})


async def force_reindex_handler(
    source_key: str,
    workspace_id: str | None = None,
    reason: str = "operator_requested",
    requested_by: str | None = None,
    metadata: dict[str, Any] | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalForceReindexResponse:
    """Handle POST /learning/retrieval/force-reindex."""
    engine = get_retrieval_engine(learning_store, agent)
    
    if engine is None:
        return RetrievalForceReindexResponse(
            request={"error": "engine not configured"}
        )
    
    try:
        if hasattr(engine, "create_force_reindex_request"):
            request = engine.create_force_reindex_request(
                workspace_id=workspace_id,
                source_key=source_key,
                reason=reason,
                requested_by=requested_by,
                metadata=metadata,
            )
            observability = get_observability_store()
            observability.record_metric(
                Metric(
                    name="adk_reindex_requests_total",
                    value=1.0,
                    unit="count",
                    attributes=_reindex_metric_attributes(request, agent),
                )
            )
            if int(request.metadata.get("coalesced_count", 0) or 0) > 0:
                observability.record_metric(
                    Metric(
                        name="adk_reindex_coalesced_total",
                        value=1.0,
                        unit="count",
                        attributes=_reindex_metric_attributes(request, agent),
                    )
                )
            _schedule_reindex_request(engine=engine, request_id=request.id, agent=agent)
            return RetrievalForceReindexResponse(request=_request_to_dict(request))
    except Exception:
        pass
    
    return RetrievalForceReindexResponse(
        request={"source_key": source_key, "status": "queued"}
    )


async def get_reindex_requests_handler(
    workspace_id: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalReindexRequestsResponse:
    """Handle GET /learning/retrieval/reindex-requests."""
    engine = get_retrieval_engine(learning_store, agent)
    
    if engine is None:
        return RetrievalReindexRequestsResponse(
            workspace_id=workspace_id,
            items=[],
        )
    
    try:
        if hasattr(engine, "list_reindex_requests"):
            requests = engine.list_reindex_requests(workspace_id=workspace_id)
            items = [
                {
                    "id": r.id,
                    "workspace_id": r.workspace_id,
                    "source_key": r.source_key,
                    "reason": r.reason,
                    "requested_by": r.requested_by,
                    "status": r.status,
                    "metadata": r.metadata,
                    "created_at": r.created_at,
                    "updated_at": r.updated_at,
                }
                for r in requests
            ]
        else:
            items = []
        
        return RetrievalReindexRequestsResponse(
            workspace_id=workspace_id,
            items=items,
        )
    except Exception:
        return RetrievalReindexRequestsResponse(
            workspace_id=workspace_id,
            items=[],
        )


async def get_reindex_request_handler(
    request_id: str,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalReindexRequestResponse:
    """Handle GET /learning/retrieval/reindex-requests/{request_id}."""
    engine = get_retrieval_engine(learning_store, agent)
    
    if engine is None:
        return RetrievalReindexRequestResponse(
            request={"error": "engine not configured"}
        )
    
    try:
        if hasattr(engine, "get_reindex_request"):
            request = engine.get_reindex_request(request_id)
            if request:
                return RetrievalReindexRequestResponse(request=_request_to_dict(request))
    except Exception:
        pass
    
    return RetrievalReindexRequestResponse(request={"error": "request not found"})


async def retry_reindex_request_handler(
    request_id: str,
    note: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalReindexRequestResponse:
    """Handle POST /learning/retrieval/reindex-requests/{request_id}/retry."""
    engine = get_retrieval_engine(learning_store, agent)
    if engine is None:
        return RetrievalReindexRequestResponse(request={"error": "engine not configured"})
    try:
        if hasattr(engine, "retry_reindex_request"):
            request = engine.retry_reindex_request(request_id, note=note)
            if request is not None:
                get_observability_store().record_metric(
                    Metric(
                        name="adk_reindex_requests_total",
                        value=1.0,
                        unit="count",
                        attributes=_reindex_metric_attributes(request, agent),
                    )
                )
                _schedule_reindex_request(engine=engine, request_id=request.id, agent=agent)
                return RetrievalReindexRequestResponse(request=_request_to_dict(request))
    except Exception:
        pass
    return RetrievalReindexRequestResponse(request={"error": "request not found"})


async def resume_reindex_request_handler(
    request_id: str,
    note: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalReindexRequestResponse:
    """Handle POST /learning/retrieval/reindex-requests/{request_id}/resume."""
    engine = get_retrieval_engine(learning_store, agent)
    if engine is None:
        return RetrievalReindexRequestResponse(request={"error": "engine not configured"})
    try:
        if hasattr(engine, "resume_reindex_request"):
            request = engine.resume_reindex_request(request_id, note=note)
            if request is not None:
                _schedule_reindex_request(engine=engine, request_id=request.id, agent=agent)
                return RetrievalReindexRequestResponse(request=_request_to_dict(request))
    except Exception:
        pass
    return RetrievalReindexRequestResponse(request={"error": "request not found"})


async def replay_reindex_request_handler(
    request_id: str,
    note: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalReindexRequestResponse:
    """Handle POST /learning/retrieval/reindex-requests/{request_id}/replay."""
    engine = get_retrieval_engine(learning_store, agent)
    if engine is None:
        return RetrievalReindexRequestResponse(request={"error": "engine not configured"})
    try:
        if hasattr(engine, "replay_reindex_request"):
            request = engine.replay_reindex_request(request_id, note=note)
            if request is not None:
                get_observability_store().record_metric(
                    Metric(
                        name="adk_reindex_requests_total",
                        value=1.0,
                        unit="count",
                        attributes=_reindex_metric_attributes(request, agent),
                    )
                )
                _schedule_reindex_request(engine=engine, request_id=request.id, agent=agent)
                return RetrievalReindexRequestResponse(request=_request_to_dict(request))
    except Exception:
        pass
    return RetrievalReindexRequestResponse(request={"error": "request not found"})


async def cancel_reindex_request_handler(
    request_id: str,
    note: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalReindexRequestResponse:
    """Handle POST /learning/retrieval/reindex-requests/{request_id}/cancel."""
    engine = get_retrieval_engine(learning_store, agent)
    if engine is None:
        return RetrievalReindexRequestResponse(request={"error": "engine not configured"})
    try:
        if hasattr(engine, "cancel_reindex_request"):
            request = engine.cancel_reindex_request(request_id, note=note)
            if request is not None:
                return RetrievalReindexRequestResponse(request=_request_to_dict(request))
    except Exception:
        pass
    return RetrievalReindexRequestResponse(request={"error": "request not found"})


async def get_reindex_summary_handler(
    workspace_id: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalOpsSummaryResponse:
    """Handle GET /learning/retrieval/summary."""
    engine = get_retrieval_engine(learning_store, agent)
    
    if engine is None:
        return RetrievalOpsSummaryResponse(
            workspace_id=workspace_id,
            requests={},
            triggers={},
            source_documents={},
            policy={},
            coalescing={},
        )
    
    try:
        if hasattr(engine, "ops_summary"):
            summary = engine.ops_summary(workspace_id=workspace_id)
            return RetrievalOpsSummaryResponse(
                workspace_id=workspace_id,
                requests=summary.get("requests", {}),
                triggers=summary.get("triggers", {}),
                source_documents=summary.get("source_documents", {}),
                policy=summary.get("policy", {}),
                coalescing=summary.get("coalescing", {}),
            )
    except Exception:
        pass
    
    return RetrievalOpsSummaryResponse(
        workspace_id=workspace_id,
        requests={},
        triggers={},
        source_documents={},
        policy={},
        coalescing={},
    )


async def get_profiles_handler(
    workspace_id: str | None = None,
    task_type: str | None = None,
    user_segment: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalProfilesResponse:
    """Handle GET /learning/retrieval/profiles."""
    engine = get_retrieval_engine(learning_store, agent)
    
    if engine is None:
        return RetrievalProfilesResponse(
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            items=[],
        )

    try:
        if hasattr(engine, "list_profiles"):
            profiles = engine.list_profiles(
                workspace_id=workspace_id,
                task_type=task_type,
                user_segment=user_segment,
            )
            items = [
                {
                    "workspace_id": p.workspace_id,
                    "task_type": p.task_type,
                    "user_segment": p.user_segment,
                    "limit": p.limit,
                    "filter": p.filter,
                    "success_count": p.success_count,
                    "failure_count": p.failure_count,
                    "last_updated": p.last_updated,
                }
                for p in profiles
            ]
        else:
            items = []

        return RetrievalProfilesResponse(
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            items=items,
        )
    except Exception:
        return RetrievalProfilesResponse(
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            items=[],
        )


async def get_source_documents_handler(
    workspace_id: str | None = None,
    source_key: str | None = None,
    learning_store: Any | None = None,
    agent: Any | None = None,
) -> RetrievalSourceDocumentsResponse:
    """Handle GET /learning/retrieval/source-documents."""
    engine = get_retrieval_engine(learning_store, agent)
    if engine is None:
        return RetrievalSourceDocumentsResponse(
            workspace_id=workspace_id,
            source_key=source_key,
            items=[],
        )
    try:
        if hasattr(engine, "list_source_documents"):
            items = []
            for document in engine.list_source_documents(
                workspace_id=workspace_id,
                source_key=source_key,
            ):
                payload = asdict(document)
                payload["content_length"] = len(document.content)
                items.append(payload)
        else:
            items = []
        return RetrievalSourceDocumentsResponse(
            workspace_id=workspace_id,
            source_key=source_key,
            items=items,
        )
    except Exception:
        return RetrievalSourceDocumentsResponse(
            workspace_id=workspace_id,
            source_key=source_key,
            items=[],
        )


__all__ = [
    "get_source_scores_handler",
    "export_source_scores_handler",
    "get_reindex_triggers_handler",
    "acknowledge_trigger_handler",
    "resolve_trigger_handler",
    "close_trigger_handler",
    "force_reindex_handler",
    "get_reindex_requests_handler",
    "get_reindex_request_handler",
    "retry_reindex_request_handler",
    "resume_reindex_request_handler",
    "replay_reindex_request_handler",
    "cancel_reindex_request_handler",
    "get_reindex_summary_handler",
    "get_profiles_handler",
    "get_source_documents_handler",
]
