from typing import Annotated
from uuid import UUID

from arthur_common.models.agent_governance_schemas import EnrichedTaskResponse
from arthur_common.models.common_schemas import PaginationParameters
from arthur_common.models.enums import PaginationSortMethod, RuleScope, RuleType
from arthur_common.models.request_schemas import (
    NewMetricRequest,
    NewRuleRequest,
    NewTaskRequest,
    SearchTasksRequest,
    UpdateMetricRequest,
    UpdateRuleRequest,
)
from arthur_common.models.response_schemas import (
    MetricResponse,
    RuleResponse,
    SearchTasksResponse,
    TaskResponse,
)
from fastapi import APIRouter, Body, Depends, HTTPException
from sqlalchemy.orm import Session
from starlette import status
from starlette.responses import RedirectResponse, Response

from clients.telemetry.telemetry_client import (
    TelemetryEventTypes,
    send_telemetry_event,
    send_telemetry_event_for_task_rule_create_completed,
)
from config.cache_config import cache_config
from dependencies import get_application_config, get_db_session, get_org_scope
from repositories.metrics_repository import MetricRepository
from repositories.rules_repository import RuleRepository
from repositories.task_polling_state_repository import TaskPollingStateRepository
from repositories.tasks_metrics_repository import TasksMetricsRepository
from repositories.tasks_repository import TaskRepository
from repositories.tasks_rules_repository import TasksRulesRepository
from routers.route_handler import GenaiEngineRoute
from routers.v2 import multi_validator
from schemas.enums import PermissionLevelsEnum
from schemas.internal_schemas import (
    ApplicationConfiguration,
    Metric,
    Rule,
    Task,
    User,
)
from utils import constants
from utils.constants import DEFAULT_ORG_ID
from utils.users import enforce_org_scope, enforce_query_org_scope, permission_checker
from utils.utils import common_pagination_parameters, public_endpoint

task_management_routes = APIRouter(
    prefix="/api/v2",
    route_class=GenaiEngineRoute,
)
rules_types = [rule.value for rule in RuleType]

################################
#### Task Management Routes ####
################################


@task_management_routes.post(
    "/tasks",
    description="Register a new task. When a new task is created, all existing default rules will be "
    "auto-applied for this new task. Optionally specify if the task is agentic.",
    response_model=TaskResponse,
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
def create_task(
    request: NewTaskRequest,
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
) -> TaskResponse:
    send_telemetry_event(TelemetryEventTypes.TASK_CREATE_INITIATED)
    if len(request.name.strip()) == 0:
        raise HTTPException(
            status_code=400,
            detail="Task names cannot contain only white space characters",
        )

    rules_repo = RuleRepository(db_session)
    tasks_repo = TaskRepository(
        db_session,
        rules_repo,
        MetricRepository(db_session),
        application_config,
    )

    # Determine the owning org from the caller's identity:
    #   admin/JWT  (org_scope is None) -> default org
    #   tenant key (org_scope is set)  -> caller's org
    if current_user is not None and current_user.org_scope is not None:
        org_id = current_user.org_scope
    else:
        org_id = DEFAULT_ORG_ID

    task = Task._from_request_model(request, org_id=org_id)
    task = tasks_repo.create_task(task)

    send_telemetry_event(TelemetryEventTypes.TASK_CREATE_COMPLETED)
    return task._to_response_model()


@task_management_routes.get(
    "/tasks",
    description="[Deprecated] Use /tasks/search endpoint. This endpoint will be removed in a future release.",
    response_model=list[TaskResponse],
    tags=["Tasks"],
    deprecated=True,
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
def get_all_tasks(
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    org_scope: UUID | None = Depends(get_org_scope),
) -> list[TaskResponse]:
    rules_repo = RuleRepository(db_session)
    tasks_repo = TaskRepository(
        db_session,
        rules_repo,
        MetricRepository(db_session),
        application_config,
    )
    # Tenant callers see only tasks in their org; admin sees everything.
    tasks = tasks_repo.get_all_tasks(org_scope=org_scope)
    return [task._to_response_model() for task in tasks]


@task_management_routes.delete(
    "/tasks/{task_id}",
    description="Archive task. Also archives all task-scoped rules. Associated default rules are unaffected.",
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
@enforce_org_scope()
def archive_task(
    task_id: UUID,
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
) -> Response:
    rules_repo = RuleRepository(db_session)
    tasks_repo = TaskRepository(
        db_session,
        rules_repo,
        MetricRepository(db_session),
        application_config,
    )
    tasks_repo.archive_task(str(task_id))
    return Response(status_code=status.HTTP_204_NO_CONTENT)


@task_management_routes.post(
    "/tasks/{task_id}/unarchive",
    description="Unarchive a previously archived task. Also unarchives all task-scoped rules and metrics that were archived with it.",
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
@enforce_org_scope()
def unarchive_task(
    task_id: UUID,
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
) -> Response:
    rules_repo = RuleRepository(db_session)
    tasks_repo = TaskRepository(
        db_session,
        rules_repo,
        MetricRepository(db_session),
        application_config,
    )
    tasks_repo.unarchive_task(str(task_id))
    return Response(status_code=status.HTTP_204_NO_CONTENT)


@task_management_routes.get(
    "/tasks/{task_id}",
    description="Get tasks.",
    response_model=TaskResponse,
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
@enforce_org_scope()
def get_task(
    task_id: UUID,
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
) -> TaskResponse:
    task_repo = TaskRepository(
        db_session,
        RuleRepository(db_session),
        MetricRepository(db_session),
        application_config,
    )
    task = task_repo.get_task_by_id(str(task_id))
    return task._to_response_model()


@task_management_routes.get(
    "/agent-tasks",
    description="Get agentic tasks with enriched agent metadata (tools, sub-agents, models). "
    "Returns only agentic tasks.",
    response_model=list[EnrichedTaskResponse],
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
def get_agent_tasks(
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    org_scope: UUID | None = Depends(get_org_scope),
) -> list[EnrichedTaskResponse]:
    """Get agentic tasks with enriched agent metadata.

    Returns tasks with additional metadata computed from spans:
    - tools: List of tools used by the agent
    - sub_agents: List of sub-agents used
    - models: List of models used
    - num_spans: Total number of spans

    Also includes creation_source information (GCP, OTEL, or manual).

    Args:
        db_session: Database session
        application_config: Application configuration
        current_user: Current authenticated user

    Returns:
        List of EnrichedTaskResponse objects (agentic tasks only)
    """
    rules_repo = RuleRepository(db_session)
    metrics_repo = MetricRepository(db_session)
    tasks_repo = TaskRepository(
        db_session,
        rules_repo,
        metrics_repo,
        application_config,
    )

    # Query only agentic tasks — tenant callers see only their own org's.
    db_tasks, _ = tasks_repo.query_tasks(
        is_agentic=True,
        include_archived=False,
        page_size=1000,  # Large page size for now, add pagination later if needed
        page=0,
        org_scope=org_scope,
    )

    # Convert to Task objects and enrich with service names
    tasks = [Task._from_database_model(db_task) for db_task in db_tasks]
    tasks = tasks_repo._enrich_tasks_with_service_names(tasks)

    # Build enriched responses
    polling_state_repo = TaskPollingStateRepository(db_session)
    enriched_responses = []
    for task in tasks:
        creation_source = tasks_repo._get_task_creation_source(task)

        # Get last_fetched from task_polling_state
        polling_state = polling_state_repo.get_by_task_id(task.id)
        last_fetched = polling_state.last_fetched if polling_state else None

        # Extract agent metadata
        agent_metadata = tasks_repo._extract_agent_metadata(task.id)

        # Convert rule links to response models
        response_rules = []
        for rule_link in task.rule_links or []:
            response_rule = rule_link.rule._to_response_model()
            response_rule.enabled = rule_link.enabled
            response_rules.append(response_rule)

        enriched_response = EnrichedTaskResponse(
            id=task.id,
            name=task.name,
            created_at=task.created_at,
            updated_at=task.updated_at,
            is_autocreated=task.is_autocreated,
            creation_source=creation_source,
            last_fetched=last_fetched,
            tools=agent_metadata["tools"],
            sub_agents=agent_metadata["sub_agents"],
            models=agent_metadata["models"],
            data_sources=agent_metadata["data_sources"],
            num_spans=agent_metadata["num_spans"],
            rules=response_rules,
        )
        enriched_responses.append(enriched_response)

    return enriched_responses


@task_management_routes.post(
    "/task",
    description="Redirect to /tasks endpoint.",
    tags=["Tasks"],
)
@public_endpoint
def redirect_to_tasks() -> RedirectResponse:
    return RedirectResponse(
        url="/api/v2/tasks",
        status_code=status.HTTP_307_TEMPORARY_REDIRECT,
    )


############################
#### Task Search Routes ####
############################


@task_management_routes.post(
    "/tasks/search",
    description="Search tasks. Can filter by task IDs, task name substring, and agentic status.",
    response_model=SearchTasksResponse,
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
@enforce_query_org_scope()
def search_tasks(
    request: SearchTasksRequest,
    pagination_parameters: Annotated[
        PaginationParameters,
        Depends(common_pagination_parameters),
    ],
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
) -> SearchTasksResponse:
    rules_repo = RuleRepository(db_session)
    metrics_repo = MetricRepository(db_session)
    tasks_repo = TaskRepository(
        db_session,
        rules_repo,
        metrics_repo,
        application_config,
    )
    db_tasks, count = tasks_repo.query_tasks(
        ids=request.task_ids,
        task_name=request.task_name,
        is_agentic=request.is_agentic,
        include_archived=getattr(request, "include_archived", False) is True,
        only_archived=getattr(request, "only_archived", False) is True,
        sort=pagination_parameters.sort or PaginationSortMethod.DESCENDING,
        page=pagination_parameters.page,
        page_size=pagination_parameters.page_size,
    )
    tasks = [Task._from_database_model(db_task) for db_task in db_tasks]
    return SearchTasksResponse(
        tasks=[task._to_response_model() for task in tasks],
        count=count,
    )


#####################################
#### Task Rule Management Routes ####
#####################################


@task_management_routes.post(
    "/tasks/{task_id}/rules",
    description="Create a rule to be applied only to this task. Available rule types are {}."
    "Note: The rules are cached by the validation endpoints for {} seconds. ".format(
        ", ".join(rules_types),
        cache_config.TASK_RULES_CACHE_TTL,
    ),
    response_model=RuleResponse,
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
@enforce_org_scope()
def create_task_rule(
    task_id: UUID,
    request: NewRuleRequest = Body(
        None,
        openapi_examples=NewRuleRequest.model_config["json_schema_extra"],  # type: ignore[arg-type]
    ),
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
) -> RuleResponse:
    send_telemetry_event(TelemetryEventTypes.TASK_RULE_CREATE_INITIATED)
    task_repo = TaskRepository(
        db_session,
        RuleRepository(db_session),
        MetricRepository(db_session),
        application_config,
    )
    rule_repo = RuleRepository(db_session)
    task = task_repo.get_task_by_id(str(task_id))

    new_rule = rule_repo.create_rule(
        Rule._from_request_model(request, scope=RuleScope.TASK),
    )
    task_repo.link_rule_to_task(task.id, new_rule.id, new_rule.type)

    response_rule = new_rule._to_response_model()
    response_rule.enabled = True
    send_telemetry_event_for_task_rule_create_completed(new_rule.type)
    return response_rule


@task_management_routes.patch(
    "/tasks/{task_id}/rules/{rule_id}",
    description="Enable or disable an existing rule for this task including the default rules.",
    response_model=TaskResponse,
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
@enforce_org_scope()
def update_task_rules(
    task_id: UUID,
    rule_id: UUID,
    body: UpdateRuleRequest,
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
) -> TaskResponse:
    task_repo = TaskRepository(
        db_session,
        RuleRepository(db_session),
        MetricRepository(db_session),
        application_config,
    )
    task_repo.toggle_task_rule_enabled(str(task_id), str(rule_id), body.enabled)
    updated_task = task_repo.get_task_by_id(str(task_id))
    return updated_task._to_response_model()


@task_management_routes.delete(
    "/tasks/{task_id}/rules/{rule_id}",
    description="Archive an existing rule for this task.",
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
@enforce_org_scope()
def archive_task_rule(
    task_id: UUID,
    rule_id: UUID,
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
) -> Response:
    rule_repo = RuleRepository(db_session)
    rule = rule_repo.get_rule_by_id(str(rule_id))

    if rule.scope == RuleScope.DEFAULT:
        raise HTTPException(
            status_code=400,
            detail=constants.ERROR_CANNOT_DELETE_DEFAULT_RULE,
        )

    tasks_rules_repo = TasksRulesRepository(db_session)
    task_rules = tasks_rules_repo._get_task_rules_ids(
        str(task_id),
        only_enabled=False,
    )
    if str(rule_id) not in task_rules:
        raise HTTPException(
            status_code=400,
            detail=constants.ERROR_UNRELATED_TASK_RULE,
        )

    task_repo = TaskRepository(
        db_session,
        rule_repo,
        MetricRepository(db_session),
        application_config,
    )

    task_repo.delete_rule_link(str(task_id), str(rule_id))

    rules, _ = rule_repo.query_rules(rule_ids=[str(rule_id)])
    rule = rules[0]
    if rule.scope == RuleScope.TASK:
        rule_repo.archive_rule(str(rule_id))

    return Response(status_code=status.HTTP_204_NO_CONTENT)


#############################
#### Task Metrics Routes ####
#############################


@task_management_routes.post(
    "/tasks/{task_id}/metrics",
    description="Create metrics for a task. Only agentic tasks can have metrics.",
    status_code=status.HTTP_201_CREATED,
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
@enforce_org_scope()
def create_task_metric(
    task_id: UUID,
    request: NewMetricRequest = Body(
        None,
        openapi_examples=NewMetricRequest.model_config["json_schema_extra"],  # type: ignore[arg-type]
    ),
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
) -> MetricResponse:
    metric_repo = MetricRepository(db_session)
    task_repo = TaskRepository(
        db_session,
        RuleRepository(db_session),
        metric_repo,
        application_config,
    )
    task = task_repo.get_task_by_id(str(task_id))
    metric = Metric._from_request_model(request)
    created_metric = metric_repo.create_metric(metric)
    task_repo.link_metric_to_task(task.id, created_metric.id)
    return created_metric._to_response_model()


@task_management_routes.patch(
    "/tasks/{task_id}/metrics/{metric_id}",
    description="Update a task metric.",
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
@enforce_org_scope()
def update_task_metric(
    task_id: UUID,
    metric_id: UUID,
    body: UpdateMetricRequest,
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
) -> TaskResponse:
    task_repo = TaskRepository(
        db_session,
        RuleRepository(db_session),
        MetricRepository(db_session),
        application_config,
    )
    task_repo.toggle_task_metric_enabled(str(task_id), str(metric_id), body.enabled)
    updated_task = task_repo.get_task_by_id(str(task_id))
    return updated_task._to_response_model()


@task_management_routes.delete(
    "/tasks/{task_id}/metrics/{metric_id}",
    description="Archive a task metric.",
    tags=["Tasks"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
@enforce_org_scope()
def archive_task_metric(
    task_id: UUID,
    metric_id: UUID,
    db_session: Session = Depends(get_db_session),
    application_config: ApplicationConfiguration = Depends(get_application_config),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
) -> Response:
    metric_repo = MetricRepository(db_session)
    task_repo = TaskRepository(
        db_session,
        RuleRepository(db_session),
        metric_repo,
        application_config,
    )
    tasks_metrics_repo = TasksMetricsRepository(db_session)

    task_metrics = tasks_metrics_repo._get_task_metrics_ids(
        str(task_id),
        only_enabled=False,
    )
    if str(metric_id) not in task_metrics:
        raise HTTPException(
            status_code=400,
            detail=constants.ERROR_UNRELATED_TASK_METRIC,
        )

    task_repo.archive_metric_link(str(task_id), str(metric_id))
    metric_repo.archive_metric(str(metric_id))

    return Response(status_code=status.HTTP_204_NO_CONTENT)
