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

import asyncio
import json
from pathlib import Path
from typing import Annotated

import typer
from rich import print as printr


def run_agent_cmd(
    file: Annotated[
        Path, typer.Argument(help='Path to a source code file to run the CL agent on')
    ],
    outfile: Annotated[
        Path,
        typer.Argument(help='Output file path to save the resulting CL Agent state'),
    ] = Path('cl_result.json'),
    artifacts: Annotated[
        bool,
        typer.Option(
            help='Should we generate artifacts like Verification Goals and Decomposition requests'
        ),
    ] = True,
):
    """
    Read the provided file and run CodeLogician Agent on it, saving the result in `outfile` file.
    """

    from imandra.u.agents.code_logician import SUPPORTED_LANGUAGES

    from codelogician.tools.cl_caller import call_cl_agent

    if not file.exists():
        typer.secho(f'🛑 File does not exist: {file}')
        raise typer.Exit(1)

    language = None
    for lang, ext in SUPPORTED_LANGUAGES:
        if file.suffix == ext:
            language = lang
            break

    if language is None:
        typer.secho(
            f'🛑 Unsupported language with extension: {file.suffix}. Only the following are supported: {SUPPORTED_LANGUAGES}'
        )
        raise typer.Exit(1)

    try:
        src_code = file.read_text()
    except Exception as e:
        typer.secho(f'🛑 Failed to read the file: {e}')
        raise typer.Exit(1)

    try:
        res = asyncio.run(
            call_cl_agent(src_code=src_code, language=language, artifacts=artifacts)
        )
    except Exception as e:
        typer.secho(
            f'🛑 Caught exception during the call: {e}', fg=typer.colors.RED, err=True
        )
        raise typer.Exit(1)

    if res is None:
        typer.secho('🛑 Failed to execute the CL agent. Received `None`', err=True)
        raise typer.Exit(1)

    printr(res)

    try:
        with open(outfile, 'w') as outfile_file:
            data = res.model_dump(mode='json')
            fields = [
                'status',
                'src_code',
                'src_lang',
                'eval_res',
                'vgs',
                'region_decomps',
                'iml_code',
                'iml_model',
            ]
            select_data = {k: data[k] for k in fields if k in data}
            print(json.dumps(select_data, indent=4), file=outfile_file)
            typer.secho(f'✅ Written results to {outfile}!')

    except Exception as e:
        typer.secho(f'🛑 Failed to print to out file: {e}', err=True)
        raise typer.Exit(1)
