from typing import Annotated

import typer
from google.cloud import aiplatform

from rag_eval.config import settings

app = typer.Typer()


@app.command()
def main(
    pipeline_spec_path: Annotated[
        str,
        typer.Option(
            "--pipeline-spec-path",
            "-p",
            help="Path to the compiled pipeline YAML file.",
        ),
    ],
    input_table: Annotated[
        str,
        typer.Option(
            "--input-table",
            "-i",
            help="Full BigQuery table name for input (e.g., 'project.dataset.table')",
        ),
    ],
    output_table: Annotated[
        str,
        typer.Option(
            "--output-table",
            "-o",
            help="Full BigQuery table name for output (e.g., 'project.dataset.table')",
        ),
    ],
    project_id: Annotated[
        str,
        typer.Option(
            "--project-id",
            help="Google Cloud project ID.",
        ),
    ] = settings.project_id,
    location: Annotated[
        str,
        typer.Option(
            "--location",
            help="Google Cloud location for the pipeline job.",
        ),
    ] = settings.location,
    display_name: Annotated[
        str,
        typer.Option(
            "--display-name",
            help="Display name for the pipeline job.",
        ),
    ] = "search-eval-pipeline-job",
):
    """Submits a Vertex AI pipeline job."""
    parameter_values = {
        "project_id": project_id,
        "location": location,
        "input_table": input_table,
        "output_table": output_table,
    }

    job = aiplatform.PipelineJob(
        display_name=display_name,
        template_path=pipeline_spec_path,
        pipeline_root=f"gs://{settings.bucket}/pipeline_root",
        parameter_values=parameter_values,
        project=project_id,
        location=location,
    )

    print(f"Submitting pipeline job with parameters: {parameter_values}")
    job.submit(
        service_account="sa-cicd-gitlab@bnt-orquestador-cognitivo-dev.iam.gserviceaccount.com"
    )
    print(f"Pipeline job submitted. You can view it at: {job._dashboard_uri()}")


if __name__ == "__main__":
    app()