import llm
import click
import contextlib
import io
import sys
import functools

@llm.hookimpl
def register_commands(cli):
    prompt_cmd = cli.commands.get('prompt')
    if not prompt_cmd:
        return

    # Check if --json is already a parameter (e.g., if the PR was merged)
    if any(p.name == "json_output" or "--json" in p.opts for p in prompt_cmd.params):
        return

    json_option = click.Option(
        ["--json"],
        is_flag=True,
        help="Output response as JSON including conversation_id"
    )
    prompt_cmd.params.append(json_option)

    original_callback = prompt_cmd.callback

    @functools.wraps(original_callback)
    def wrapped_prompt(*args, **kwargs):
        json_output = kwargs.pop('json', False)
        
        if not json_output:
            return original_callback(*args, **kwargs)

        no_log = kwargs.get('no_log')
        if no_log:
            raise click.ClickException("Cannot use --json with --no-log as the response must be logged to output JSON")

        captured_responses = []
        original_log_to_db = llm.Response.log_to_db
        original_async_log_to_db = llm.AsyncResponse.log_to_db

        def patched_log_to_db(self, db):
            captured_responses.append(self)
            return original_log_to_db(self, db)
            
        def patched_async_log_to_db(self, db):
            captured_responses.append(self)
            return original_async_log_to_db(self, db)

        llm.Response.log_to_db = patched_log_to_db
        llm.AsyncResponse.log_to_db = patched_async_log_to_db
        
        dummy_stdout = io.StringIO()
        try:
            with contextlib.redirect_stdout(dummy_stdout):
                original_callback(*args, **kwargs)
        finally:
            llm.Response.log_to_db = original_log_to_db
            llm.AsyncResponse.log_to_db = original_async_log_to_db

        if captured_responses:
            from llm.cli import logs_list
            ctx = click.get_current_context()
            
            database = kwargs.get('database')
            if database:
                import pathlib
                log_path = pathlib.Path(database)
            else:
                from llm.cli import logs_db_path
                log_path = logs_db_path()

            # We need all the unique response IDs that were captured
            response_ids = [resp.id for resp in captured_responses]
            if response_ids:
                # `logs_list` handles fetching entries.
                # Just invoke `logs_list` with the first response ID
                ctx.invoke(
                    logs_list,
                    count=len(response_ids),
                    database=str(log_path),
                    id_gte=response_ids[0],
                    json_output=True,
                )

    prompt_cmd.callback = wrapped_prompt
