import time
import typing as T
import json
import uuid
from pathlib import Path

import httpx
import typer
from rich.console import Console
from rich.table import Table
from rich.progress import track
from pydantic import BaseModel, TypeAdapter

err_console = Console(stderr=True)
console = Console()


app = typer.Typer(no_args_is_help=True)


def valid_input(schema: dict, input: dict) -> bool:
    # TODO
    if False:
        console.print("[bold red]Error: invalid schema!")
        return False
    return True


class StepResponseError(BaseModel):
    exception_log: str | None = None
    raw_response_text: str | None = None


class StepResponse(BaseModel):
    # it must always return json? Make this a requirement
    body: dict | None
    headers: dict | None

    error: StepResponseError | None = None


class StepRun(BaseModel):
    id: str
    step_id: str

    execute_response: StepResponse
    execute_duration_ms: float

    started_at_ms: float
    ended_at_ms: float


class Step(BaseModel):
    id: str
    slug: str
    plugin_id: str
    input: dict[str, T.Any]


class Plugin(BaseModel):
    input_schema: dict[str, T.Any]
    execute_url: str


class StepsInput(BaseModel):
    steps: list[Step]
    plugins: dict[str, Plugin]


@app.command()
def run(
    *,
    json_file: T.Annotated[
        T.Optional[str], typer.Option(help="Path to the JSON file to run.")
    ] = None,
    json_string: T.Annotated[
        T.Optional[str], typer.Option(help="JSON string to run.")
    ] = None,
    output_file: T.Annotated[str, typer.Option(help="Path to the output file.")],
    output_length: T.Annotated[
        T.Optional[int],
        typer.Option(help="Number of characters to print for the output of each step."),
    ] = 200,
):
    """
    Provide either a --json_file or a --json_string to run.
    Provide an optional --output-length for the length of the output of the run to display.
    """
    if not json_file and not json_string:
        err_console.print(
            "[bold red]Error:[/bold red] Please provide either [green]--json_file[/green] or [green]--json_string[/green]."
        )
        return
    if json_file and json_string:
        err_console.print(
            "[bold red]Error:[/bold red] Please provide either [green]--json_file[/green] or [green]--json_string[/green], not both."
        )
        return
    if json_file:
        path = Path(json_file)
        if path.exists():
            with path.open("r") as f:
                json_string = f.read()
        else:
            err_console.print(
                f"[bold red]Error:[/bold red] No file found at {json_file}"
            )
            return
    steps_input = StepsInput.model_validate_json(json_string)
    runs: list[StepRun] = []
    for step in track(steps_input.steps, description="Running steps..."):
        console.log(f"running step {step.id}")
        step_started_at_ms = time.time() * 1_000
        # confirm the input types conform to the plugins inputSchema
        input = step.input
        plugin = steps_input.plugins[step.plugin_id]
        input_schema = plugin.input_schema
        valid_input(schema=input_schema, input=input)
        run_id = f"run_{uuid.uuid4()}"
        r = httpx.post(
            url=plugin.execute_url,
            json=input,
            headers={
                "x-whitecloak-run-id": str(run_id),
                "x-whitecloak-step-id": step.id,
            },
            timeout=60,
        )
        step_ended_at_ms = time.time() * 1_000
        try:
            step_response = StepResponse(body=r.json(), headers=r.headers)
        except Exception as e:
            step_response = StepResponse(
                body=None,
                headers=None,
                error=StepResponseError(exception_log=str(e), raw_response_text=r.text),
            )
        runs.append(
            StepRun(
                id=run_id,
                step_id=step.id,
                execute_response=step_response,
                execute_duration_ms=r.elapsed.total_seconds() * 1_000,
                started_at_ms=step_started_at_ms,
                ended_at_ms=step_ended_at_ms,
            )
        )
    table = Table("Run Id", "Step Id", "Latency (ms)", "Result")
    for _run in runs:
        table.add_row(
            _run.step_id,
            _run.step_id,
            str(round(_run.execute_duration_ms, 2)),
            json.dumps(_run.execute_response.body)[0:output_length],
        )
    console.print(table)
    path = Path(output_file)
    if not path.exists():
        # do this check just in case you do not have filesystem directory write access
        path.parent.mkdir(parents=True, exist_ok=True)
        v = TypeAdapter(list[StepRun]).dump_json(runs)
        path.write_bytes(v)


@app.command()
def hi():
    console.print("hi")


if __name__ == "__main__":
    app()
