from typing import Annotated, Optional
from uuid import UUID, uuid4

from arthur_common.models.common_schemas import PaginationParameters
from fastapi import APIRouter, Depends, HTTPException, Path, Query, Response, status
from sqlalchemy.orm import Session

from dependencies import get_db_session, get_org_scope, get_validated_task
from repositories.prompt_experiment_repository import PromptExperimentRepository
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.prompt_experiment_schemas import (
    CreatePromptExperimentRequest,
    PromptExperimentDetail,
    PromptExperimentListResponse,
    PromptExperimentSummary,
    PromptVersionResultListResponse,
    TestCaseListResponse,
)
from services.prompt_experiment_executor import PromptExperimentExecutor
from utils.users import enforce_org_scope, permission_checker
from utils.utils import common_pagination_parameters

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


@prompt_experiment_routes.get(
    "/tasks/{task_id}/prompt_experiments",
    summary="List prompt experiments",
    description="List all prompt experiments for a task with optional filtering and pagination",
    response_model=PromptExperimentListResponse,
    response_model_exclude_none=True,
    tags=["Prompt Experiments"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
@enforce_org_scope()
def list_prompt_experiments(
    pagination_parameters: Annotated[
        PaginationParameters,
        Depends(common_pagination_parameters),
    ],
    search: Optional[str] = Query(
        None,
        description="Search text to filter experiments by name, description, prompt name, or dataset name",
    ),
    dataset_id: Optional[UUID] = Query(
        None,
        description="Filter experiments by dataset ID",
    ),
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    task: Task = Depends(get_validated_task),
) -> PromptExperimentListResponse:
    """
    List all prompt experiments for a given task.

    Returns paginated list of experiment summaries.
    Optionally filter by search text matching experiment name, description, prompt name, or dataset name.
    Optionally filter by dataset ID.
    """
    try:
        repo = PromptExperimentRepository(db_session)
        experiments, total_count = repo.list_experiments(
            task_id=task.id,
            pagination_params=pagination_parameters,
            search_text=search,
            dataset_id=dataset_id,
        )

        page = pagination_parameters.page
        page_size = pagination_parameters.page_size
        total_pages = (
            (total_count + page_size - 1) // page_size if total_count > 0 else 0
        )

        return PromptExperimentListResponse(
            data=experiments,
            page=page,
            page_size=page_size,
            total_pages=total_pages,
            total_count=total_count,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db_session.close()


@prompt_experiment_routes.post(
    "/tasks/{task_id}/prompt_experiments",
    summary="Create and run a prompt experiment",
    description="Create a new prompt experiment and initiate execution",
    response_model=PromptExperimentSummary,
    response_model_exclude_none=True,
    status_code=status.HTTP_200_OK,
    tags=["Prompt Experiments"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
@enforce_org_scope()
def create_prompt_experiment(
    experiment_request: CreatePromptExperimentRequest,
    db_session: Session = Depends(get_db_session),
    current_user: User | None = Depends(multi_validator.validate_api_multi_auth),
    task: Task = Depends(get_validated_task),
) -> PromptExperimentSummary:
    """
    Create a new prompt experiment and start execution.

    The experiment will test the specified prompt configurations (saved or unsaved)
    against the dataset using the configured evaluations.
    """
    try:
        # Validate at least one prompt is provided
        if (
            not experiment_request.prompt_configs
            or len(experiment_request.prompt_configs) == 0
        ):
            raise HTTPException(
                status_code=400,
                detail="At least one prompt configuration is required",
            )

        # Check for duplicate prompt configs
        prompt_keys_seen = set()
        for config in experiment_request.prompt_configs:
            if config.type == "saved":
                key = f"saved:{config.name}:{config.version}"
            else:  # unsaved
                # For unsaved prompts, we check for identical message/model combinations
                # since auto_name will be generated by the backend
                key = f"unsaved:{hash(str(config.messages))}"

            if key in prompt_keys_seen:
                raise HTTPException(
                    status_code=400,
                    detail=f"Duplicate prompt configuration detected: {key}",
                )
            prompt_keys_seen.add(key)

        # Generate a unique experiment ID
        experiment_id = str(uuid4())

        repo = PromptExperimentRepository(db_session)
        experiment = repo.create_experiment(
            task_id=task.id,
            experiment_id=experiment_id,
            request=experiment_request,
        )

        # Kick off async execution of the experiment
        executor = PromptExperimentExecutor()
        executor.execute_experiment_async(experiment_id)

        return experiment
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db_session.close()


@prompt_experiment_routes.get(
    "/prompt_experiments/{experiment_id}",
    summary="Get prompt experiment details",
    description="Get detailed information about a specific prompt experiment including summary results",
    response_model=PromptExperimentDetail,
    response_model_exclude_none=True,
    tags=["Prompt Experiments"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
def get_prompt_experiment(
    experiment_id: str = Path(
        ...,
        description="The ID of the experiment to retrieve",
        title="Experiment ID",
    ),
    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),
) -> PromptExperimentDetail:
    """
    Get detailed information about a prompt experiment.

    Returns full experiment configuration and summary results across all test cases.
    """
    try:
        repo = PromptExperimentRepository(db_session)
        return repo.get_experiment(experiment_id, org_scope=org_scope)
    except HTTPException:
        raise
    except ValueError as e:
        if "not found" in str(e).lower():
            raise HTTPException(status_code=404, detail=str(e))
        else:
            raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db_session.close()


@prompt_experiment_routes.get(
    "/prompt_experiments/{experiment_id}/test_cases",
    summary="Get experiment test cases",
    description="Get paginated list of test case results for a prompt experiment",
    response_model=TestCaseListResponse,
    response_model_exclude_none=True,
    tags=["Prompt Experiments"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
def get_experiment_test_cases(
    pagination_parameters: Annotated[
        PaginationParameters,
        Depends(common_pagination_parameters),
    ],
    experiment_id: str = Path(
        ...,
        description="The ID of the experiment",
        title="Experiment ID",
    ),
    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),
) -> TestCaseListResponse:
    """
    Get detailed test case results for an experiment.

    Returns paginated list of individual test cases with their inputs, outputs,
    and evaluation results.
    """
    try:
        repo = PromptExperimentRepository(db_session)
        test_cases, total_count = repo.get_test_cases(
            experiment_id=experiment_id,
            pagination_params=pagination_parameters,
            org_scope=org_scope,
        )

        page = pagination_parameters.page
        page_size = pagination_parameters.page_size
        total_pages = (
            (total_count + page_size - 1) // page_size if total_count > 0 else 0
        )

        return TestCaseListResponse(
            data=test_cases,
            page=page,
            page_size=page_size,
            total_pages=total_pages,
            total_count=total_count,
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db_session.close()


@prompt_experiment_routes.get(
    "/prompt_experiments/{experiment_id}/prompts/{prompt_key}/results",
    summary="Get prompt results",
    description="Get paginated list of results for a specific prompt within an experiment (supports both saved and unsaved prompts)",
    response_model=PromptVersionResultListResponse,
    response_model_exclude_none=True,
    tags=["Prompt Experiments"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_READ.value)
def get_prompt_version_results(
    pagination_parameters: Annotated[
        PaginationParameters,
        Depends(common_pagination_parameters),
    ],
    experiment_id: str = Path(
        ...,
        description="The ID of the experiment",
        title="Experiment ID",
    ),
    prompt_key: str = Path(
        ...,
        description="The prompt key (format: 'saved:name:version' or 'unsaved:auto_name'). URL-encode colons as %3A",
        title="Prompt Key",
    ),
    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),
) -> PromptVersionResultListResponse:
    """
    Get detailed results for a specific prompt within an experiment.

    Returns paginated list of results showing inputs, rendered prompt, outputs,
    and evaluation results for the specified prompt across all test cases.

    The prompt_key parameter should be:
    - For saved prompts: 'saved:name:version' (e.g., 'saved:my_prompt:1')
    - For unsaved prompts: 'unsaved:auto_name' (e.g., 'unsaved:unsaved_prompt_1')

    Note: Colons in the prompt_key should be URL-encoded as %3A in the URL path.
    """
    try:
        repo = PromptExperimentRepository(db_session)
        results, total_count = repo.get_prompt_version_results(
            experiment_id=experiment_id,
            prompt_key=prompt_key,
            pagination_params=pagination_parameters,
            org_scope=org_scope,
        )

        page = pagination_parameters.page
        page_size = pagination_parameters.page_size
        total_pages = (
            (total_count + page_size - 1) // page_size if total_count > 0 else 0
        )

        return PromptVersionResultListResponse(
            data=results,
            page=page,
            page_size=page_size,
            total_pages=total_pages,
            total_count=total_count,
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db_session.close()


@prompt_experiment_routes.delete(
    "/prompt_experiments/{experiment_id}",
    summary="Delete prompt experiment",
    description="Delete a prompt experiment and all its associated data",
    status_code=status.HTTP_204_NO_CONTENT,
    responses={
        status.HTTP_204_NO_CONTENT: {"description": "Experiment deleted successfully."},
    },
    tags=["Prompt Experiments"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
def delete_prompt_experiment(
    experiment_id: str = Path(
        ...,
        description="The ID of the experiment to delete",
        title="Experiment ID",
    ),
    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:
    """
    Delete a prompt experiment.

    This will remove the experiment and all associated test case results.
    This operation cannot be undone.
    """
    try:
        repo = PromptExperimentRepository(db_session)
        repo.delete_experiment(experiment_id, org_scope=org_scope)
        return Response(status_code=status.HTTP_204_NO_CONTENT)
    except HTTPException:
        raise
    except ValueError as e:
        if "not found" in str(e).lower():
            raise HTTPException(status_code=404, detail=str(e))
        else:
            raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db_session.close()


@prompt_experiment_routes.patch(
    "/prompt_experiments/{experiment_id}/notebook",
    summary="Attach notebook to experiment",
    description="Attach a notebook to an existing experiment",
    response_model=PromptExperimentSummary,
    response_model_exclude_none=True,
    tags=["Prompt Experiments"],
)
@permission_checker(permissions=PermissionLevelsEnum.TASK_WRITE.value)
def attach_notebook_to_experiment(
    experiment_id: str = Path(..., description="ID of the experiment"),
    notebook_id: str = Query(..., description="ID of the notebook to attach"),
    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),
) -> PromptExperimentSummary:
    """Attach a notebook to an existing experiment."""
    try:
        repo = PromptExperimentRepository(db_session)
        return repo.attach_notebook_to_experiment(
            experiment_id, notebook_id, org_scope=org_scope
        )
    except HTTPException:
        raise
    except ValueError as e:
        if "not found" in str(e).lower():
            raise HTTPException(status_code=404, detail=str(e))
        else:
            raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db_session.close()
