import json
import logging
import sys
from typing import List
from urllib.parse import urlparse

import click

from highlighter.agent.agent import _validate_uuid
from highlighter.client import CloudAgent
from highlighter.client.agents import create_agent_token, create_machine_agent_version
from highlighter.client.base_models.base_models import UserType
from highlighter.client.gql_client import HLClient
from highlighter.client.json_tools import HLJSONEncoder
from highlighter.core.config import HighlighterRuntimeConfigError, _parse_header_pairs
from highlighter.core.runtime import Runtime, derive_agent_name

logger = logging.getLogger(__name__)


@click.group("agent")
@click.pass_context
def agent_group(ctx):
    pass


@agent_group.command("start")
@click.option("--name", "-n", type=str, default=None, help="Override agent runtime name")
@click.option("--task-id", "-t", type=str, default=None, multiple=True, help="comma separate for multiple")
@click.option("--step-id", "-i", type=str, default=None)
@click.option("--stream-definitions-file", type=str)
@click.option("--dump-definition", type=str, default=None)
@click.option("--allow-non-machine-user", is_flag=True, default=False)
@click.option("--server", is_flag=True, default=False)
@click.option("--telemetry", is_flag=True, default=False, help="Enable telemetry")
@click.option("--telemetry-rate", type=int, default=None, show_default="config or 1")
@click.option(
    "--telemetry-otlp-endpoint",
    type=str,
    default=None,
    envvar=["HL_TELEMETRY_OTLP_ENDPOINT", "OTEL_EXPORTER_OTLP_ENDPOINT"],
)
@click.option(
    "--telemetry-otlp-headers",
    type=str,
    default=None,
)
@click.option(
    "--telemetry-otlp-protocol",
    type=str,
    default=None,
    help="OTLP protocol: 'http' or 'grpc' (defaults to config value).",
)
@click.argument("agent_definition", type=click.Path(dir_okay=False, exists=False))
@click.argument("inputs", nargs=-1, type=click.STRING, required=False)
@click.pass_context
def _start(
    ctx,
    name,
    task_id,
    step_id,
    stream_definitions_file,
    dump_definition,
    allow_non_machine_user,
    server,
    telemetry,
    telemetry_rate,
    telemetry_otlp_endpoint,
    telemetry_otlp_headers,
    telemetry_otlp_protocol,
    agent_definition,
    inputs,
):
    """Start a local Highlighter Agent to process data either from your local machine or from Highlighter tasks.

    When processing URLs or local files, a single stream is created to process each URL or file.
    The Agent definition must have its first element as a
    DataSourceCapability, such as ImageDataSource, VideoDataSource,
    TextDataSource, JsonArrayDataSource, etc. The examples below assume
    the use of ImageDataSource.

    When processing Highlighter tasks, a single stream is created for each
    task. The Agent definition should use AssessmentRead as the first element in this case.
    Note: When processing tasks, use a GraphQL API key specific to the agent being run. You can
    create this using 'hl agent create-token'.

    URLs and files can be passed as arguments.
    Alterlatively, the argument '-' can be given and the contents of an image
    can be passed via stdin (when using the ImageDataSource capability) or JSON can be
    piped in (when using the JsonArrayDataSource capability).

    If '--server' is specified, the runtime won't shutdown after processing the specified
    inputs.

    Examples:

      \b
      1. Start an agent against a single image path
      \b
        > hl agent start agent-def.json images/123.jpg

      \b
      2. Start an agent against a multiple image paths
      \b
        > hl agent start agent-def.json $(find images/ -name *.jpg)
        OR
        > find images/ -name *.jpg -print0 | xargs -0 hl agent start agent-def.json

      \b
      3. Cat the contents of an image to an agent
      \b
        > cat images/123.jpg | hl agent start agent-def.json -

      \b
      4. Pipe json data to stdin
      \b
        > cat '[{"foo": "bar"},{"foo": "baz"}]' | hl agent start agent-def.json -

      \b
      5. Process streams defined in a file
      \b
        > hl agent start agent-def.json --stream-definitions-file streams.json

      \b
      6. Process tasks from a Highlighter machine-step, using a local agent definition
      \b
        > STEP_UUID=xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
        > hl agent start agent-def.json --step-id "$STEP_UUID"

      \b
      7. Process tasks from a Highlighter machine-step, using an agent definition from Highlighter
      \b
        > STEP_UUID=xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
        > AGENT_UUID=xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
        > hl agent start "$AGENT_UUID" --step-id "$STEP_UUID"

    """
    try:
        stream_definitions = []
        files = []
        urls = []
        if len(inputs) > 0:
            if inputs[0] == "-":
                urls.append("hlpipe://")
            else:
                for input_arg in inputs:
                    if urlparse(input_arg).scheme == "":
                        files.append(input_arg)
                    else:
                        urls.append(input_arg)
        if stream_definitions_file is not None:
            with open(stream_definitions_file) as sdf:
                stream_definitions = json.load(sdf)
        telemetry_profiler = None
        telemetry_cfg = None
        if ctx.obj.get("hl_cfg") is not None:
            telemetry_cfg = ctx.obj["hl_cfg"].telemetry
        telemetry_enabled = telemetry_cfg.enabled if telemetry_cfg is not None else False
        if telemetry:
            telemetry_enabled = True
        if telemetry_rate is None:
            if telemetry_cfg is not None:
                telemetry_rate = telemetry_cfg.sample_rate
            else:
                telemetry_rate = 1
        agent_name = derive_agent_name(agent_definition, name)

        if not allow_non_machine_user:
            if agent_version_id := _validate_uuid(agent_definition):
                client = HLClient.get_client()
                current_user = client.current_user(return_type=UserType)
                if current_user.machine_agent_version_id is None:
                    logger.info(f"Creating temporary machine token for agent version {agent_version_id}")
                    token = create_agent_token(str(agent_version_id))
                    os.environ["HL_WEB_GRAPHQL_API_TOKEN"] = token
                    # Reset the client instance so it reloads with the new token
                    HLClient.clear_instance()
                    # Also need to reset it in ctx if it was stored there
                    if ctx.obj.get("client"):
                        ctx.obj["client"] = HLClient.get_client()

        if telemetry_enabled:
            from highlighter.core.telemetry import Profiler, attach_otel_log_handler

            telemetry_otlp_endpoint_resolved = telemetry_otlp_endpoint
            telemetry_otlp_protocol_resolved = telemetry_otlp_protocol
            headers = telemetry_cfg.resolved_headers() if telemetry_cfg is not None else {}
            if telemetry_otlp_headers:
                cli_headers = _parse_header_pairs(telemetry_otlp_headers)
                headers.update(cli_headers)
            if telemetry_otlp_endpoint_resolved is None and telemetry_cfg is not None:
                telemetry_otlp_endpoint_resolved = telemetry_cfg.otlp_endpoint
            if telemetry_otlp_protocol_resolved is None and telemetry_cfg is not None:
                telemetry_otlp_protocol_resolved = telemetry_cfg.otlp_protocol
            if not telemetry_otlp_endpoint_resolved:
                raise click.ClickException(
                    "Telemetry enabled but no OTLP endpoint configured.\n"
                    "Set one of:\n"
                    "  --telemetry-otlp-endpoint <url>\n"
                    "  HL_TELEMETRY_OTLP_ENDPOINT=<url>\n"
                    "  OTEL_EXPORTER_OTLP_ENDPOINT=<url>\n"
                    "  [telemetry] otlp_endpoint in ~/.highlighter/config"
                )
            resource_attributes: dict[str, str] = {}
            if agent_name:
                resource_attributes["highlighter.agent_name"] = agent_name
                resource_attributes["service.name"] = agent_name
            telemetry_profiler = Profiler(
                run_id=str(ctx.obj.get("run_id") or "cli"),
                sample_rate=telemetry_rate,
                otlp_endpoint=telemetry_otlp_endpoint_resolved,
                otlp_headers=headers,
                otlp_protocol=telemetry_otlp_protocol_resolved,
                resource_attributes=resource_attributes,
            )
            attach_otel_log_handler(
                otlp_endpoint=telemetry_otlp_endpoint_resolved,
                otlp_protocol=telemetry_otlp_protocol_resolved or "http",
                headers=headers,
                resource_attributes=telemetry_profiler.resource_attributes_for_otel(),
                logger=logging.getLogger(),
            )

        runtime = Runtime(
            agent_definition=agent_definition,
            dump_definition=dump_definition,
            allow_non_machine_user=allow_non_machine_user,
            hl_cfg=ctx.obj.get("hl_cfg"),
            hl_client=ctx.obj.get("client"),
            queue_response=ctx.obj.get("queue_response"),
            name=agent_name,
            telemetry_profiler=telemetry_profiler,
        )
        runtime.run(
            stream_definitions=stream_definitions,
            files=files,
            urls=urls,
            task_ids=task_id,
            step_id=step_id,
            server=server,
        )
    except click.ClickException:
        raise
    except HighlighterRuntimeConfigError as e:
        logger.error(f"Failed to load configuration: {e}")
        sys.exit(2)
    except Exception as exc:
        logger.error(f"{type(exc).__qualname__}: {exc}")
        import traceback

        traceback.print_exc()
        sys.exit(1)


@agent_group.command("create-token")
@click.option("--machine-agent-version-id", type=str)
def _create_token(
    machine_agent_version_id,
):
    """Create an access token for an agent

    Once an access token has been created, run an agent with that identity by
    setting `HL_WEB_GRAPHQL_API_TOKEN=<new-token>` before running `hl agent start`

    Provide the ID of the machine-agent-version in Highlighter.
    """
    machine_agent_token = create_agent_token(machine_agent_version_id)
    # Print to stdout rather than log
    print(machine_agent_token)


@agent_group.command("check")
@click.argument("agent_definition", type=click.Path(dir_okay=False, exists=True))
@click.option(
    "--cloud",
    is_flag=True,
    help="Validate referenced UUIDs against Highlighter Cloud",
)
@click.option(
    "--verbose",
    is_flag=True,
    help="Print all validated UUID references",
)
@click.pass_context
def _check(ctx, agent_definition, cloud, verbose):
    """Validate a pipeline definition without running it."""
    from highlighter.agent.agent import load_pipeline_definition
    from highlighter.agent.pipeline_validator import (
        validate_pipeline_definition_cloud,
        validate_pipeline_definition_structure,
    )

    try:
        pipeline_def = load_pipeline_definition(agent_definition)

        print(f"Validating pipeline: {pipeline_def.get('name', 'Unknown')}")
        print(f"File: {agent_definition}")
        print()

        print("[1/2] Validating structure...")
        structure_result = validate_pipeline_definition_structure(pipeline_def)

        if not structure_result.is_valid:
            print(f"  ✗ FAILED with {len(structure_result.errors)} error(s):")
            for error in structure_result.errors:
                print(f"\n  Error: {error.message}")
                if error.location:
                    print(f"  Location: {error.location}")
                if error.context:
                    print(f"  Context: {error.context}")
            sys.exit(1)

        print("  ✓ Structure validation passed")

        if structure_result.warnings:
            print(f"\n  ⚠ {len(structure_result.warnings)} warning(s):")
            for warning in structure_result.warnings:
                print(f"    - {warning.message}")

        if cloud:
            client = None
            if ctx.obj:
                client = ctx.obj.get("client")

            if not hasattr(client, "endpoint_url"):
                print("\n✗ Cloud validation requires an authenticated HLClient.")
                sys.exit(1)

            print("\n[2/2] Validating cloud references...")
            print(f"Cloud endpoint: {client.endpoint_url}")
            print("Cloud auth: api-token")

            cloud_result, details = validate_pipeline_definition_cloud(pipeline_def, client)
            if verbose:
                print("\nCloud checks:")
                for check in details.checks:
                    status = "ok" if check.get("ok") else "fail"
                    extra = ""
                    if check.get("error"):
                        extra = f" error={check['error']}"
                    print(
                        " -"
                        f" path={check['path']}"
                        f" type={check['entity_type']}"
                        f" query={check['query']}"
                        f" arg={check['arg']}"
                        f" uuid={check['uuid']}"
                        f" status={status}"
                        f"{extra}"
                    )
            if not cloud_result.is_valid:
                print("\nCloud reference failures:")
                for issue in details.invalid + details.unsupported + details.missing:
                    hint = issue.hint or "not found"
                    print(f"- path={issue.path} type={issue.entity_type} uuid={issue.uuid} hint={hint}")

                print("\nChecks summary:")
                print(f"Validated UUIDs: {len(details.validated)}")
                print(f"Missing UUIDs: {len(details.missing) + len(details.unsupported)}")
                print(f"Invalid UUIDs: {len(details.invalid)}")
                print("Cache: disabled")
                sys.exit(1)

            print("\nValidated references:")
            for ref in sorted(details.validated, key=lambda item: (item.path, item.entity_type, item.uuid)):
                print(f"- path={ref.path} type={ref.entity_type} uuid={ref.uuid}")

            print("\nChecks summary:")
            print(f"Validated UUIDs: {len(details.validated)}")
            print(f"Missing UUIDs: {len(details.missing)}")
            print(f"Invalid UUIDs: {len(details.invalid)}")
            print("Cache: disabled")

        print("\n" + "=" * 60)
        print("VALIDATION SUMMARY")
        print("=" * 60)
        print(f"Pipeline: {pipeline_def.get('name', 'Unknown')}")
        print("Status: ✓ READY FOR USE")
        print(f"Elements: {len(pipeline_def.get('elements', []))}")
        print(f"Graph stages: {len(pipeline_def.get('graph', []))}")
        print("=" * 60)

    except click.ClickException:
        raise
    except Exception as e:
        print(f"\n✗ Validation failed with an unexpected error: {e}")
        sys.exit(1)


@agent_group.command("list")
@click.option("--cloud", is_flag=True)
def _list(cloud):
    """List running agents"""
    if not cloud:
        raise NotImplementedError(
            "Can currently only list agents running in the cloud with the '--cloud' flag"
        )
    try:
        client = HLClient.get_client()
        cloud_agents = client.cloudAgents(return_type=List[CloudAgent])
        print(json.dumps(cloud_agents, indent=4, cls=HLJSONEncoder))
    except Exception as e:
        print(f"Error listing agents. {type(e).__qualname__}: {e}")
