#
# Imandra Inc.
#
# endpts_sketches.py
#

import asyncio
import logging

from fastapi import HTTPException

from codelogician.sketch.sketch import (
    SketchChange,
    SketchChangeResult,
    SketchChangeTask,
)
from codelogician.sketch.task import (
    SketchChgDeleteDef,
    SketchChgInsertDef,
    SketchChgModifyDef,
    SketchChgSetModel,
)
from codelogician.strategy.events import (
    SketchChangeEvent,
    SketchChangeResultEvent,
)
from codelogician.strategy.state import StrategyState
from codelogician.strategy.worker import run_sketch_task

from ..cl_server import CLServer

log = logging.getLogger(__name__)


def register_sketches_endpoints(app: CLServer):  # noqa: C901
    """Register sketch-related endpoints"""

    @app.get('/sketches/search', operation_id='search_sketches')
    async def search(query: str) -> list:
        """
        Search for specific sketch
        """
        return app.search_sketches(query or '')

    @app.get('/sketches/list', operation_id='list_sketches')
    def get_sketches_list() -> dict[str, list[dict[str, str]]]:
        """
        List all of the sketches available the currently selected strategy.
        Returns a list of dictionaries, references by strategy id and mapped to a list of sketch IDs.

        'sketch_id': sketch_id,
        'anchor_model': self.sketches[sketch_id].anchor_model_path

        Example: [
            {'sketch_id': "SketchID1", 'anchormodel': "path1.py"},
            {'sketch_id': "SketchID2", 'anchromodel': "path2.py"}
        """

        return app.list_sketches()

    @app.post('/sketches/create', operation_id='create_sketch')
    def create_sketch(anchor_model_path: str):
        """
        Create a new sketch for this strategy
        """

        strategy_id = app.current_strat_id
        if strategy_id is None:
            raise HTTPException(status_code=404, detail='No current strategy')

        try:
            strat_state: StrategyState = app.strategy_state_by_id(strategy_id)
        except Exception as e:
            raise HTTPException(
                status_code=404, detail=f'Failed to get strategy state: {e}'
            )

        if anchor_model_path not in strat_state.curr_meta_model.models:
            raise HTTPException(
                status_code=404, detail=f'Unknown model path: {anchor_model_path}'
            )

        try:
            sketch_id = app.create_sketch(strategy_id, anchor_model_path)
        except Exception as e:
            raise HTTPException(
                status_code=404, detail=f'Failed to create a new sketch:{e}'
            )

        return sketch_id

    @app.post('/sketches/{sketch_id}/try_change', operation_id='try_sketch_change')
    async def try_sketch_change(
        sketch_id: str, change: SketchChange, commit_on_success: bool = True
    ) -> SketchChangeResult | str:
        """
        Try a sketch change (optionally and if successful, apply it - by default, this is True)
        """

        sketch = app.get_sketch_from_sketch_id(sketch_id)

        if sketch is None:
            raise HTTPException(
                status_code=404, detail=f'Could not locate sketch with ID={sketch_id}'
            )

        try:
            result_iml = sketch.process_change(sketch_id, change)
        except Exception as e:
            raise HTTPException(status_code=404, detail=f'Could not apply change: {e}')

        try:
            change_result = asyncio.run(
                run_sketch_task(
                    SketchChangeTask(
                        sketch_id=sketch_id, change=change, iml_code=result_iml
                    )
                )
            )
        except Exception as e:
            raise HTTPException(
                status_code=404, detail=f'Error during call to ImandraX: {e}'
            )

        if (
            commit_on_success
            and change_result is not None
            and change_result.error is None
        ):
            strat = app.get_strat_for_sketch_id(sketch_id=sketch_id)
            if strat:
                strat.add_event(
                    SketchChangeResultEvent(
                        sketch_id=sketch_id, change_result=change_result
                    )
                )

                log.info(
                    f'Added SketchChangeResult to strategy[id={strat.state().strat_id}]'
                )
            else:
                log.error(f'Failed to find strategy for sketch_id = {sketch_id}')

        if change_result is None:
            return 'N//A'
        else:
            return change_result

    @app.post('/sketches/{sketch_id}/change', operation_id='apply_sketch_change')
    def apply_sketch_change(
        sketch_id: str,
        change: SketchChgSetModel
        | SketchChgInsertDef
        | SketchChgModifyDef
        | SketchChgDeleteDef,
    ):
        """
        Apply sketch change
        """

        if sketch_id not in app.all_sketch_ids():
            raise HTTPException(
                status_code=404, detail=f'Sketch with id={sketch_id} was not found!'
            )

        strat = app.get_strat_for_sketch_id(sketch_id)

        if strat is None:
            raise HTTPException(
                status_code=400,
                detail=f'Could not locate strategy for sketch={sketch_id}',
            )

        strat.add_event(SketchChangeEvent(sketch_id=sketch_id, change=change))

        return 'OK'

    @app.post('/sketches/{sketch_id}/rollback', operation_id='rollback_changes')
    def rollback_changes(sketch_id: str, target_state_id: int | None = None):
        """
        Rollback changes. If `target_state_id` is specified, then the Sketch will go back to the specified state ID if possible.
        If `target_state_id` is not specified, then the Sketch will rollback the last change, unless it's the initial one.
        """

        sketch = app.get_sketch_from_sketch_id(sketch_id)

        if sketch is None:
            raise HTTPException(
                status_code=404, detail=f'Sketch with id={sketch_id} was not found!'
            )

        try:
            sketch.rollback(target_state_id=target_state_id)
        except Exception as e:
            raise HTTPException(
                status_code=403, detail=f'Failed to rollback the state: {e}'
            )

        return 'OK'

    @app.post('/sketches/{sketch_id}/state', operation_id='')
    def get_latest_sketch_state(sketch_id: str):
        """
        Return the latest sketch state
        """

        if sketch_id not in app.all_sketch_ids():
            raise HTTPException(
                status_code=404, detail=f'Could not locate sketch with id = {sketch_id}'
            )

        return 'OK'
