# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_prism

from pathlib import Path

import click
from loguru import logger

from coreason_identity.models import UserContext
from coreason_identity.types import SecretStr
from coreason_prism.interface import Prism, PrismMode


@click.group()  # type: ignore[misc]
@click.option("--light-mode", is_flag=True, help="Enable Light Mode (skip heavy models).")  # type: ignore[misc]
@click.pass_context  # type: ignore[misc]
def cli(ctx: click.Context, light_mode: bool) -> None:
    """Coreason Prism: The Scientific Eye / Multi-Modal Encoder.

    Initializes the Prism context and handles global options.

    Args:
        ctx (click.Context): The Click context object.
        light_mode (bool): If True, enables Light Mode to skip loading heavy models.

    Raises:
        click.ClickException: If initialization fails.
    """
    ctx.ensure_object(dict)
    # Initialize Prism with the global light_mode flag
    # We store it in the context object to be reused or init here
    # Since commands might need specific arguments, we can init Prism here and store in context
    try:
        ctx.obj["prism"] = Prism(light_mode=light_mode)
    except Exception as e:
        logger.exception("Failed to initialize Prism")
        raise click.ClickException(f"Initialization failed: {e}") from e


@cli.command()  # type: ignore[misc]
@click.argument("input_data")  # type: ignore[misc]
@click.pass_context  # type: ignore[misc]
def molecule(ctx: click.Context, input_data: str) -> None:
    """Process a molecule from SMILES, InChI, SELFIES string or file path.

    Outputs the result as a JSON string.

    Args:
        ctx (click.Context): The Click context object containing the Prism instance.
        input_data (str): The molecule string or file path.

    Raises:
        click.ClickException: If molecule processing fails.
    """
    prism: Prism = ctx.obj["prism"]
    system_context = UserContext(
        user_id=SecretStr("cli-user"),
        roles=["system"],
        metadata={"source": "cli"},
    )
    try:
        # Check if input is a file path
        path_obj = Path(input_data)
        if path_obj.exists():
            result = prism.process_molecule(path_obj, context=system_context)
        else:
            result = prism.process_molecule(input_data, context=system_context)

        # Output result as JSON
        click.echo(result.model_dump_json(indent=2))

    except Exception as e:
        logger.error(f"Molecule processing failed: {e}")
        raise click.ClickException(str(e)) from e


@cli.command()  # type: ignore[misc]
@click.argument("image_path", type=click.Path(exists=True))  # type: ignore[misc]
@click.argument("source_id")  # type: ignore[misc]
@click.option(  # type: ignore[misc]
    "--mode",
    type=click.Choice([m.value for m in PrismMode], case_sensitive=False),
    required=True,
    help="Processing mode: CHART or BIO",
)
@click.pass_context  # type: ignore[misc]
def image(ctx: click.Context, image_path: str, source_id: str, mode: str) -> None:
    """Process an image (Chart or Bio-Image).

    Outputs the extracted data and metadata as a JSON string.

    Args:
        ctx (click.Context): The Click context object containing the Prism instance.
        image_path (str): Path to the image file.
        source_id (str): Identifier for the source document.
        mode (str): Processing mode ('CHART' or 'BIO').

    Raises:
        click.ClickException: If image processing fails.
    """
    prism: Prism = ctx.obj["prism"]
    system_context = UserContext(
        user_id=SecretStr("cli-user"),
        roles=["system"],
        metadata={"source": "cli"},
    )
    try:
        # Convert string mode to Enum
        prism_mode = PrismMode(mode.upper())

        result = prism.process_image(Path(image_path), source_id, prism_mode, context=system_context)

        # Output result as JSON
        click.echo(result.model_dump_json(indent=2))

    except Exception as e:
        logger.error(f"Image processing failed: {e}")
        raise click.ClickException(str(e)) from e


if __name__ == "__main__":
    cli()  # pragma: no cover
