from typing import Annotated
from uuid import UUID

from arthur_common.models.common_schemas import PaginationParameters
from arthur_common.models.task_eval_schemas import (
    ListTraceTransformsResponse,
    TraceTransformResponse,
)
from fastapi import APIRouter, Depends, HTTPException, Path
from sqlalchemy.orm import Session
from starlette.responses import Response
from starlette.status import HTTP_204_NO_CONTENT

from dependencies import (
    get_db_session,
    get_org_scope,
    get_validated_task,
    transform_list_filter_parameters,
)
from repositories.metrics_repository import MetricRepository
from repositories.span_repository import SpanRepository
from repositories.tasks_metrics_repository import TasksMetricsRepository
from repositories.trace_transform_repository import TraceTransformRepository
from routers.route_handler import GenaiEngineRoute
from routers.v2 import multi_validator
from schemas.enums import PermissionLevelsEnum
from schemas.internal_schemas import Task, User
from schemas.request_schemas import (
    NewTraceTransformRequest,
    TraceTransformUpdateRequest,
    TransformListFilterRequest,
)
from schemas.response_schemas import (
    ListTraceTransformVersionsResponse,
    TraceTransformVersionResponse,
    TransformDependents,
    TransformExtractionResponseList,
)
from utils.transform_executor import execute_transform
from utils.users import enforce_org_scope, permission_checker
from utils.utils import common_pagination_parameters

transform_routes = APIRouter(
    prefix="/api/v1",
    route_class=GenaiEngineRoute,
)


@transform_routes.get(
    "/tasks/{task_id}/traces/transforms",
    description="List all transforms for a task.",
    response_model=ListTraceTransformsResponse,
    tags=["Transforms"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
@enforce_org_scope()
def list_transforms_for_task(
    pagination_parameters: Annotated[
        PaginationParameters,
        Depends(common_pagination_parameters),
    ],
    filter_request: Annotated[
        TransformListFilterRequest,
        Depends(transform_list_filter_parameters),
    ],
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    task: Task = Depends(get_validated_task),
) -> ListTraceTransformsResponse:
    try:
        trace_transform_repo = TraceTransformRepository(db_session)
        transforms = trace_transform_repo.list_transforms(
            task.id,
            pagination_parameters,
            filter_request,
        )
        definitions = trace_transform_repo.get_latest_definitions_for_transforms(
            [t.id for t in transforms],
        )
        return ListTraceTransformsResponse(
            transforms=[
                transform.to_response_model(definition=definitions[transform.id])
                for transform in transforms
            ],
            count=len(transforms),
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@transform_routes.get(
    "/traces/transforms/{transform_id}",
    description="Get a specific transform.",
    response_model=TraceTransformResponse,
    tags=["Transforms"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
def get_transform(
    transform_id: UUID = Path(description="ID of the transform to fetch."),
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    org_scope: UUID | None = Depends(get_org_scope),
) -> TraceTransformResponse:
    try:
        trace_transform_repo = TraceTransformRepository(db_session)
        trace_transform = trace_transform_repo.get_transform_by_id(
            transform_id, org_scope=org_scope
        )

        if not trace_transform:
            raise HTTPException(
                status_code=404,
                detail=f"Transform {transform_id} not found",
            )

        definition = trace_transform_repo.get_latest_definition(transform_id)
        return trace_transform.to_response_model(definition=definition)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@transform_routes.get(
    "/traces/transforms/{transform_id}/dependents",
    description="Get resources that depend on this transform.",
    response_model=TransformDependents,
    tags=["Transforms"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
def get_transform_dependents(
    transform_id: UUID = Path(description="ID of the transform."),
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    org_scope: UUID | None = Depends(get_org_scope),
) -> TransformDependents:
    try:
        repo = TraceTransformRepository(db_session)
        if not repo.get_transform_by_id(transform_id, org_scope=org_scope):
            raise HTTPException(
                status_code=404,
                detail=f"Transform {transform_id} not found",
            )
        return repo.get_transform_dependents(transform_id)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@transform_routes.post(
    "/tasks/{task_id}/traces/transforms",
    description="Create a new transform for a task.",
    response_model=TraceTransformResponse,
    tags=["Transforms"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
@enforce_org_scope()
def create_transform_for_task(
    request: NewTraceTransformRequest,
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    task: Task = Depends(get_validated_task),
) -> TraceTransformResponse:
    try:
        trace_transform_repo = TraceTransformRepository(db_session)
        trace_transform = trace_transform_repo.create_transform(task.id, request)
        return trace_transform.to_response_model(definition=request.definition)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@transform_routes.patch(
    "/traces/transforms/{transform_id}",
    description="Update a transform.",
    response_model=TraceTransformResponse,
    tags=["Transforms"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
def update_transform(
    request: TraceTransformUpdateRequest,
    transform_id: UUID = Path(description="ID of the transform to update."),
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    org_scope: UUID | None = Depends(get_org_scope),
) -> TraceTransformResponse:
    try:
        trace_transform_repo = TraceTransformRepository(db_session)
        trace_transform = trace_transform_repo.update_transform(
            transform_id, request, org_scope=org_scope
        )
        definition = trace_transform_repo.get_latest_definition(transform_id)
        return trace_transform.to_response_model(definition=definition)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@transform_routes.get(
    "/traces/transforms/{transform_id}/versions",
    description="List all version snapshots for a transform, ordered by version number descending.",
    response_model=ListTraceTransformVersionsResponse,
    tags=["Transforms"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
def list_transform_versions(
    transform_id: UUID = Path(description="ID of the transform."),
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    org_scope: UUID | None = Depends(get_org_scope),
) -> ListTraceTransformVersionsResponse:
    try:
        repo = TraceTransformRepository(db_session)
        if not repo.get_transform_by_id(transform_id, org_scope=org_scope):
            raise HTTPException(
                status_code=404,
                detail=f"Transform {transform_id} not found",
            )
        return repo.list_versions(transform_id)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@transform_routes.get(
    "/traces/transforms/{transform_id}/versions/{version_id}",
    description="Get a specific version snapshot of a transform.",
    response_model=TraceTransformVersionResponse,
    tags=["Transforms"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
def get_transform_version(
    transform_id: UUID = Path(description="ID of the transform."),
    version_id: UUID = Path(description="ID of the version to fetch."),
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    org_scope: UUID | None = Depends(get_org_scope),
) -> TraceTransformVersionResponse:
    try:
        repo = TraceTransformRepository(db_session)
        if not repo.get_transform_by_id(transform_id, org_scope=org_scope):
            raise HTTPException(
                status_code=404,
                detail=f"Transform {transform_id} not found",
            )
        return repo.get_version_by_id(transform_id, version_id)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@transform_routes.delete(
    "/traces/transforms/{transform_id}/versions/{version_id}",
    description="Delete a specific version of a transform. Returns 409 if it is the only version or is pinned by a continuous eval.",
    tags=["Transforms"],
    status_code=HTTP_204_NO_CONTENT,
    responses={
        409: {
            "description": "Version is the only one remaining or is pinned by a continuous eval.",
        },
    },
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
def delete_transform_version(
    transform_id: UUID = Path(description="ID of the transform."),
    version_id: UUID = Path(description="ID of the version to delete."),
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    org_scope: UUID | None = Depends(get_org_scope),
) -> Response:
    try:
        repo = TraceTransformRepository(db_session)
        if not repo.get_transform_by_id(transform_id, org_scope=org_scope):
            raise HTTPException(
                status_code=404,
                detail=f"Transform {transform_id} not found",
            )
        repo.delete_version(transform_id, version_id)
        return Response(status_code=HTTP_204_NO_CONTENT)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@transform_routes.delete(
    "/traces/transforms/{transform_id}",
    description="Delete a transform. Returns 409 if the transform is referenced by continuous evals, agentic experiments, or agentic notebooks.",
    tags=["Transforms"],
    status_code=HTTP_204_NO_CONTENT,
    responses={
        409: {
            "description": "Transform has dependent resources that must be removed first.",
        },
    },
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
def delete_transform(
    transform_id: UUID = Path(description="ID of the transform to delete."),
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    org_scope: UUID | None = Depends(get_org_scope),
) -> Response:
    try:
        trace_transform_repo = TraceTransformRepository(db_session)
        trace_transform_repo.delete_transform(transform_id, org_scope=org_scope)
        return Response(status_code=HTTP_204_NO_CONTENT)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@transform_routes.post(
    "/traces/{trace_id}/transforms/{transform_id}/extractions",
    description="Execute a transform against a trace to extract variables.",
    response_model=TransformExtractionResponseList,
    tags=["Transforms"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
def execute_trace_transform_extraction(
    trace_id: str = Path(
        description="ID of the trace to execute the transform against.",
    ),
    transform_id: UUID = Path(description="ID of the transform to execute."),
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    org_scope: UUID | None = Depends(get_org_scope),
) -> TransformExtractionResponseList:
    try:
        # Fetch the transform (org-scoped for tenant callers)
        trace_transform_repo = TraceTransformRepository(db_session)
        trace_transform = trace_transform_repo.get_transform_by_id(
            transform_id,
            org_scope=org_scope,
        )

        if not trace_transform:
            raise HTTPException(
                status_code=404,
                detail=f"Transform {transform_id} not found",
            )

        # Fetch the trace (org-scoped for tenant callers)
        tasks_metrics_repo = TasksMetricsRepository(db_session)
        metrics_repo = MetricRepository(db_session)

        span_repo = SpanRepository(db_session, tasks_metrics_repo, metrics_repo)

        trace = span_repo.get_trace_by_id(
            trace_id=trace_id,
            include_metrics=False,
            compute_new_metrics=False,
            org_scope=org_scope,
        )

        if not trace:
            raise HTTPException(
                status_code=404,
                detail=f"Trace with ID {trace_id} not found",
            )

        # Fetch the latest version to get the transform definition
        versions = trace_transform_repo.list_versions(transform_id).versions
        if not versions:
            raise HTTPException(
                status_code=404,
                detail=f"No versions found for transform {transform_id}",
            )

        # Execute the transform
        return execute_transform(
            trace=trace,
            transform_definition=versions[0].definition,
        )

    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
