#
#   Imandra Inc.
#
#   codelogician/commands/server/cmd_model.py
#

from __future__ import annotations

import json
import urllib.parse
from typing import Any

import httpx
import typer

app = typer.Typer(name='model', help='CLI for CodeLogician model endpoints.')


# -------------------------
# Helpers
# -------------------------


def _print_json(data: Any) -> None:
    typer.echo(json.dumps(data, indent=2, ensure_ascii=False))


def _handle_http_error(resp: httpx.Response) -> None:
    try:
        payload = resp.json()
    except Exception:
        payload = None

    if isinstance(payload, dict) and 'detail' in payload:
        raise typer.BadParameter(f'HTTP {resp.status_code}: {payload["detail"]}')
    raise typer.BadParameter(f'HTTP {resp.status_code}: {resp.text}')


def _get(
    base_url: str,
    path: str,
    *,
    params: dict[str, Any] | None = None,
    timeout: float,
) -> Any:
    url = base_url.rstrip('/') + path
    with httpx.Client(timeout=timeout) as client:
        resp = client.get(url, params=params)
    if resp.status_code >= 400:
        _handle_http_error(resp)
    return resp.json()


def _post(
    base_url: str,
    path: str,
    *,
    json_body: Any | None = None,
    timeout: float,
) -> Any:
    url = base_url.rstrip('/') + path
    with httpx.Client(timeout=timeout) as client:
        resp = client.post(url, json=json_body)
    if resp.status_code >= 400:
        _handle_http_error(resp)

    ctype = resp.headers.get('content-type', '')
    if 'application/json' in ctype:
        return resp.json()
    return resp.text


def _cfg(ctx: typer.Context) -> dict[str, Any]:
    root = ctx.find_root()
    assert isinstance(root.obj, dict)
    return root.obj


# -------------------------
# Read endpoints
# -------------------------


@app.command('paths')
def model_paths(ctx: typer.Context) -> None:
    """GET /model/paths — list all model relative paths."""
    cfg = _cfg(ctx)
    data = _get(cfg['base_url'], '/model/paths', timeout=cfg['timeout'])

    if cfg['json_out']:
        _print_json(data)
    else:
        for p in data:
            typer.echo(p)


@app.command('get')
def model_get(
    ctx: typer.Context,
    index: int | None = typer.Option(None, '--index', help='Model index (0-based).'),
    path: str | None = typer.Option(None, '--path', help='Model relative path.'),
) -> None:
    """
    GET a model by index or path.

    - GET /model/byindex/{index}
    - GET /model/bypath/{path}
    """
    cfg = _cfg(ctx)

    if (index is None) == (path is None):
        raise typer.BadParameter('Specify exactly one of --index or --path')

    if index is not None:
        raw = _get(cfg['base_url'], f'/model/byindex/{index}', timeout=cfg['timeout'])
        # Endpoint may return JSON (dict) or JSON-encoded string; normalize to python object.
        if isinstance(raw, str):
            try:
                data: Any = json.loads(raw)
            except Exception:
                typer.echo(raw)
                return
        else:
            data = raw
    else:
        enc_path = urllib.parse.quote(path or '', safe='')
        data = _get(
            cfg['base_url'], f'/model/bypath/{enc_path}', timeout=cfg['timeout']
        )

    if cfg['json_out']:
        _print_json(data)
    else:
        typer.echo(data)


# -------------------------
# Command endpoints
# -------------------------


@app.command('freeze')
def model_freeze(
    ctx: typer.Context,
    index: int = typer.Argument(..., help='Model index (0-based).'),
) -> None:
    """POST /model/cmd/freeze/{index} — freeze a model's IML code."""
    cfg = _cfg(ctx)
    out = _post(cfg['base_url'], f'/model/cmd/freeze/{index}', timeout=cfg['timeout'])
    typer.echo(out if isinstance(out, str) else json.dumps(out))


@app.command('unfreeze')
def model_unfreeze(
    ctx: typer.Context,
    index: int = typer.Argument(..., help='Model index (0-based).'),
) -> None:
    """POST /model/cmd/unfreeze/{index} — unfreeze a model's IML code."""
    cfg = _cfg(ctx)
    out = _post(cfg['base_url'], f'/model/cmd/unfreeze/{index}', timeout=cfg['timeout'])
    typer.echo(out if isinstance(out, str) else json.dumps(out))


@app.command('cmd')
def model_command(
    ctx: typer.Context,
    rel_path: str = typer.Argument(
        ..., help='Model relative path (as returned by `paths`).'
    ),
    cmd_json: str = typer.Option(
        ...,
        '--cmd-json',
        help='JSON object for the Command payload (sent as request body).',
    ),
) -> None:
    """
    POST /model/cmd/{path} — submit a CodeLogician agent command for a model.

    Example:
      --cmd-json '{'kind':'AutoFormalize'}'
    """
    cfg = _cfg(ctx)

    try:
        cmd_obj = json.loads(cmd_json)
    except json.JSONDecodeError as e:
        raise typer.BadParameter(f'--cmd-json must be valid JSON: {e}') from e

    enc_path = urllib.parse.quote(rel_path, safe='')
    out = _post(
        cfg['base_url'],
        f'/model/cmd/{enc_path}',
        json_body=cmd_obj,
        timeout=cfg['timeout'],
    )
    typer.echo(out if isinstance(out, str) else json.dumps(out, indent=2))


if __name__ == '__main__':
    app()
