#
#   Imandra Inc.
#
#   codelogician/commands/cmd_eval.py
#

# pyright: strict
# pyright: reportMissingTypeStubs=false

import asyncio
import json as jsonlib
import os
import sys
from functools import partial
from pathlib import Path
from typing import Annotated, Any, Literal, TypedDict, assert_never, cast

import structlog
import typer
from imandrax_api.twirp.exceptions import TwirpServerException
from imandrax_api_models import DecomposeRes, InstanceRes, Task, VerifyRes
from imandrax_api_models.client import (
    ImandraXAsyncClient,
    ImandraXClient,
    get_imandrax_async_client,
    get_imandrax_client,
)
from imandrax_api_models.context_utils import (
    format_decomp_res,
    format_eval_res,
    format_vg_res,
    remove_art_and_task_fields,
)
from imandrax_api_models.logging_utils import configure_logging
from imandrax_api_models.pp_goal_state import PO_RES_PP_BIN_PATH, pp_goal_state
from iml_query.multifile import gather_modules, parse_imports
from iml_query.processing import (
    get_decomp_reqs,
    get_instance_reqs,
    get_verify_reqs,
)
from iml_query.processing.decomp import DecompReqArgs
from pydantic import TypeAdapter

from codelogician.util import (
    ImandraAPIAccessDenied,
    ImandraAPIKeyError,
    get_imandra_uni_key,
)

logger = structlog.getLogger(__name__)

app = typer.Typer(name='ImandraX')

HELP_ENV_VAR = """\
Environment variables:
    IMANDRAX_ENV: (dev, prod) environment for ImandraX
    CODELOGICIAN_DEBUG: (true, false) debug mode
"""

_get_imandrax_client = partial(get_imandrax_client, auth_token=get_imandra_uni_key())
_get_imandrax_async_client = partial(
    get_imandrax_async_client, auth_token=get_imandra_uni_key()
)

if os.environ.get('CODELOGICIAN_DEBUG', 'false').lower() == 'true':
    configure_logging('debug')
else:
    configure_logging('warning')


def _remove_fields(data: dict[str, Any]) -> dict[str, Any]:
    """Resursively look inside a dict for certain keys and remove it."""
    data = data.copy()
    remove_fields = ['artifact', 'task']
    for k in list(data.keys()):
        v = data[k]
        if k in remove_fields:
            data.pop(k)
        elif isinstance(v, dict):
            v = cast(dict[str, Any], v)
            data[k] = _remove_fields(v)
    return data


def _load_iml(path: str | None) -> str:
    """Read IML code from a file or stdin."""
    if (path is None) or (path == '-'):
        # Read from stdin if no path is provided
        iml = sys.stdin.read()
        if parse_imports(iml):
            msg = (
                'Error: IML content passed via stdin contains `[@@@import ...]` statements, which are ignored.'
                'Pass the IML file to properly load imported modules '
            )
            logger.error(msg)
            typer.secho(msg, err=True)
            typer.Exit(1)
    elif not Path(path).exists():
        raise typer.BadParameter(f'IML file {path} does not exist')
    else:
        iml = gather_modules(Path(path))
    return iml


def ifprintf(message: Any | None) -> None:
    """Similar to OCmal's `ifprintf` function.

    Used to intercept stdout printing when "--json" is set.
    """
    pass


# Eval
# ====================


@app.command(
    name='check',
    help='Evaluate an IML file without VG and decomp.',
)
def check(
    file: Annotated[
        str | None,
        typer.Argument(
            help='Path of the IML file to evaluate. Set to "-" to read from stdin.',
        ),
    ] = None,
    with_vgs: Annotated[
        bool,
        typer.Option(
            help='Whether to strip verify and instance requests before evaluating.',
        ),
    ] = False,
    with_decomps: Annotated[
        bool,
        typer.Option(
            help='Whether to decomp requests before evaluating.',
        ),
    ] = False,
    json: Annotated[
        bool,
        typer.Option(
            help='Whether to output the results in JSON format.',
        ),
    ] = False,
):
    iml = _load_iml(file)

    try:
        c: ImandraXClient = _get_imandrax_client()
        eval_res = c.eval_model(src=iml, with_vgs=with_vgs, with_decomps=with_decomps)

        goal_state_s: str | None = None
        # If eval success, and has PO res, and PO res is no-proof,
        # we try to get the subgoals
        if eval_res.success:
            if PO_RES_PP_BIN_PATH is not None and eval_res.tasks:
                task: Task = eval_res.tasks[
                    0
                ]  # TODO(Q): is this always the first task?
                art_list_res: list[str] = c.list_artifacts(task).kinds
                if 'po_res' in art_list_res:
                    try:
                        po_res_art_zip = c.get_artifact_zip(task, kind='po_res')
                        goal_state_s = pp_goal_state(po_res_art_zip)
                        logger.debug('Got goal state', goal_state_s=goal_state_s)
                    except Exception as e:
                        logger.error('Failed to get goal state', exc_info=str(e))
            else:
                logger.info('PO_RES_PP_BIN_PATH is not set')

        # Return
        if not json:
            # str
            out = format_eval_res(eval_res, iml)
            if goal_state_s:
                out += f'\n\n- goal state:\n{goal_state_s}'
            typer.echo(out)
        else:
            # json str
            out = eval_res.model_dump(mode='json')
            out['goal_state'] = goal_state_s
            typer.echo(jsonlib.dumps(out, indent=2))
    except TwirpServerException as e:
        reason: str = cast(str, e.meta['body']['error'])  # pyright: ignore[reportUnknownMemberType]
        match cast(str, e.message):
            case 'Unauthorized':
                typer.secho(
                    f'Error: Access to Imandra API not authorized ({reason}).\n'
                    'Hint: Check you have a valid API key in your IMANDRA_UNI_KEY environment variable.',
                    err=True,
                )
            case 'Forbidden':
                typer.secho(
                    f'Error: Access to Imandra API forbidden ({reason}).\n'
                    f'Hint: {ImandraAPIAccessDenied("").hint}',
                    err=True,
                )
            case _:
                raise
    except ImandraAPIKeyError as e:
        typer.secho(f'Error: {str(e)}\nHint: {e.hint}', err=True)


# VG
# ====================


class VGItem(TypedDict):
    kind: Literal['verify', 'instance']
    src: str
    hints: str | None
    start_point: tuple[int, int]
    end_point: tuple[int, int]


def _collect_vgs(iml: str) -> list[VGItem]:
    iml_, verify_reqs, verify_req_ranges = get_verify_reqs(iml)
    _iml, instance_reqs, instance_req_ranges = get_instance_reqs(iml_)

    # Collect
    vg_items: list[VGItem] = []
    for req, req_range in zip(verify_reqs, verify_req_ranges, strict=True):
        vg_items.append(
            {
                'kind': 'verify',
                'src': req['src'],
                'hints': req['hints'],
                'start_point': (req_range.start_point[0], req_range.start_point[1]),
                'end_point': (req_range.end_point[0], req_range.end_point[1]),
            }
        )
    for req, req_range in zip(instance_reqs, instance_req_ranges, strict=False):
        vg_items.append(
            {
                'kind': 'instance',
                'src': req['src'],
                'hints': req['hints'],
                'start_point': (req_range.start_point[0], req_range.start_point[1]),
                'end_point': (req_range.end_point[0], req_range.end_point[1]),
            }
        )
    vg_items.sort(key=lambda x: x['start_point'])
    return vg_items


@app.command(
    name='list-vg',
    help='List verification goals specified by `verify` or `instance` commands in an IML file.',
)
def list_vg(
    file: Annotated[
        str | None,
        typer.Argument(
            help='Path of the IML file to check. Set to "-" to read from stdin.',
        ),
    ] = None,
    json: Annotated[
        bool,
        typer.Option(
            help='Whether to output the results in JSON format.',
        ),
    ] = False,
):
    iml = _load_iml(file)

    vgs: list[VGItem] = _collect_vgs(iml)

    if not json:
        for i, item in enumerate(vgs, 1):
            loc_str = (
                f'{item["start_point"][0]}:{item["start_point"][1]}'
                f'-{item["end_point"][0]}:{item["end_point"][1]}'
            )
            hint_str = f' {item["hints"]}' if item['hints'] else ''
            typer.echo(f'{i}: {item["kind"]} ({loc_str}): {item["src"]}{hint_str}')
    else:
        json_s = TypeAdapter(list[VGItem]).dump_json(vgs, indent=2)
        typer.echo(json_s)


@app.command(
    name='check-vg',
    help='Check verification goals specified by `verify` or `instance` commands in an IML file.',
)
def check_vg(
    file: Annotated[
        str | None,
        typer.Argument(
            help='Path of the IML file to check. Set to "-" to read from stdin.',
        ),
    ] = None,
    index: Annotated[
        list[int] | None,
        typer.Option(
            help='Name of the verification goal to check.',
        ),
    ] = None,
    check_all: Annotated[
        bool,
        typer.Option(
            help='Whether to check all verify requests in the IML file.',
        ),
    ] = False,
    json: Annotated[
        bool,
        typer.Option(
            help='Whether to output the results in JSON format.',
        ),
    ] = False,
):
    index = index or []

    if not json:
        echo = typer.echo
    else:
        echo = ifprintf

    async def _async_check_vg() -> list[VerifyRes | InstanceRes]:
        iml = _load_iml(file)
        vgs = _collect_vgs(iml)

        index_: list[int] = (
            list(range(1, len(vgs) + 1)) if check_all or (len(index) == 0) else index
        )

        vg_with_idx: list[tuple[int, VGItem]] = [
            (i, vg) for (i, vg) in enumerate(vgs, 1) if i in index_
        ]

        async def _check_vg(
            vg: VGItem,
            i: int,
            c: ImandraXAsyncClient,
        ) -> VerifyRes | InstanceRes:
            match vg['kind']:
                case 'verify':
                    res = await c.verify_src(src=vg['src'], hints=vg['hints'])
                case 'instance':
                    res = await c.instance_src(src=vg['src'], hints=vg['hints'])
                case _:
                    assert_never(vg['kind'])
            echo(f'{i}: {vg["kind"]} ({vg["src"]})')
            echo(format_vg_res(res))
            return res

        async with _get_imandrax_async_client() as c:
            eval_res = await c.eval_model(src=iml)
            echo(format_eval_res(eval_res, iml))
            if eval_res.has_errors:
                echo('Error(s) found in IML file. Exiting.')
                sys.exit(1)
                return
            echo('\n' + '=' * 5 + 'VG' + '=' * 5 + '\n')
            tasks = [_check_vg(vg, i, c) for (i, vg) in vg_with_idx]
            return await asyncio.gather(*tasks)

    vg_res_list: list[VerifyRes | InstanceRes] = asyncio.run(_async_check_vg())
    if json:
        out: list[dict[str, Any]] = [
            remove_art_and_task_fields(vg_res.model_dump(mode='json'))
            for vg_res in vg_res_list
        ]
        typer.echo(jsonlib.dumps(out, indent=2))


# decomp
# ====================


class DecompItem(TypedDict):
    req_args: DecompReqArgs
    start_point: tuple[int, int]
    end_point: tuple[int, int]


def _collect_decomps(iml: str) -> list[DecompItem]:
    iml, decomp_reqs, ranges = get_decomp_reqs(iml)

    decomp_items: list[DecompItem] = [
        DecompItem(
            req_args=req,
            start_point=range_.start_point,
            end_point=range_.end_point,
        )
        for req, range_ in zip(decomp_reqs, ranges, strict=True)
    ]

    decomp_items.sort(key=lambda x: x['start_point'])
    return decomp_items


@app.command(
    name='list-decomp',
    help='List decomp requests in an IML file.',
)
def list_decomp(
    file: Annotated[
        str | None,
        typer.Argument(
            help='Path of the IML file to check. Set to "-" to read from stdin.',
        ),
    ] = None,
    json: Annotated[
        bool,
        typer.Option(
            help='Whether to output the results in JSON format.',
        ),
    ] = False,
):
    iml = _load_iml(file)
    decomps: list[DecompItem] = _collect_decomps(iml)

    if not json:
        for i, item in enumerate(decomps, 1):
            typer.echo(f'{i}: {item["req_args"]["name"]}')
    else:
        json_s = TypeAdapter(list[DecompItem]).dump_json(decomps, indent=2)
        typer.echo(json_s)


@app.command(
    name='check-decomp',
    help='Check decomp requests in an IML file.',
)
def check_decomp(
    file: Annotated[
        str | None,
        typer.Argument(
            help='Path of the IML file to check. Set to "-" to read from stdin.',
        ),
    ] = None,
    index: Annotated[
        list[int] | None,
        typer.Option(
            help='Index of the decomposition request to check.',
        ),
    ] = None,
    check_all: Annotated[
        bool,
        typer.Option(
            help='Whether to check all decomp requests in the IML file.',
        ),
    ] = False,
    json: Annotated[
        bool,
        typer.Option(
            help='Whether to output the results in JSON format.',
        ),
    ] = False,
):
    index = index or []

    if not json:
        echo = typer.echo
    else:
        echo = ifprintf

    async def _async_check_decomp() -> list[DecomposeRes]:
        iml = _load_iml(file)
        decomps = _collect_decomps(iml)

        index_: list[int] = (
            list(range(1, len(decomps) + 1))
            if check_all or (len(index) == 0)
            else index
        )

        decomp_with_idx: list[tuple[int, DecompItem]] = [
            (i, decomp) for (i, decomp) in enumerate(decomps, 1) if i in index_
        ]

        async def _check_decomp(
            decomp: DecompItem, i: int, c: ImandraXAsyncClient
        ) -> DecomposeRes:
            res = await c.decompose(**decomp['req_args'])
            echo(f'{i}: decompose {decomp["req_args"]["name"]}')
            echo(format_decomp_res(res))
            return res

        async with _get_imandrax_async_client() as c:
            eval_res = await c.eval_model(src=iml)
            echo(format_eval_res(eval_res, iml))
            if eval_res.has_errors:
                typer.echo('Error(s) found in IML file. Exiting.')
                sys.exit(1)
                return

            echo('\n' + '=' * 5 + 'Decomp' + '=' * 5 + '\n')
            tasks = [_check_decomp(decomp, i, c) for (i, decomp) in decomp_with_idx]
            return await asyncio.gather(*tasks)

    decomp_res_list: list[DecomposeRes] = asyncio.run(_async_check_decomp())

    if json:
        out: list[dict[str, Any]] = [
            remove_art_and_task_fields(decomp_res.model_dump(mode='json'))
            for decomp_res in decomp_res_list
        ]
        typer.echo(jsonlib.dumps(out, indent=2))


# Conditionally add gen-test command
try:
    from imandrax_codegen.__main__ import (
        gen_test_command,
    )

    if not os.environ.get('IMANDRAX_API_KEY'):
        # gen-test needs IMANDRAX_API_KEY
        os.environ['IMANDRAX_API_KEY'] = get_imandra_uni_key()

    app.command(
        name='gen-test',
        help='Generate test cases in Python / TypeScript from region decomp in IML.',
    )(gen_test_command)
except ImandraAPIKeyError:
    # gen-test command unavailable without API key
    pass
except ValueError as e:
    if 'Only MacOS and Linux are supported' in str(e):
        pass
    else:
        raise e


if __name__ == '__main__':
    app()
