#
#   Imandra Inc.
#
#   codelogician/sketch/sketch.py
#

import uuid
from pathlib import Path

from imandra.u.agents.code_logician.base import FormalizationStatus
from imandra.u.agents.code_logician.base.region_decomp import RegionDecomp
from imandra.u.agents.code_logician.base.vg import VG
from pydantic import BaseModel, Field

from codelogician.util import find, maybe, maybe_else

from ..modeling.model import Model
from ..tools.iml_utils import (
    add_definition,
    extract_definitions,
    remove_definition,
    topo_sort_definitions,
    update_definition,
)
from .task import (
    SketchChange,
    SketchChangeResult,
    SketchChangeTask,
    SketchChgDeleteDef,
    SketchChgInsertDef,
    SketchChgModifyDef,
    SketchChgSetModel,
)


def _iml_path(rel_path: str) -> str:
    return str(Path(rel_path).with_suffix('.iml'))


def _norm_def(body: str) -> str:
    return body.strip() + '\n'


def _ensure_trailing_nl(s: str) -> str:
    s2 = s.rstrip()
    return (s2 + '\n') if s2 else ''


class VGChange(BaseModel):
    """
    Helps us track VG changes over different states
    """

    is_new: bool  # Is this a new VG that was just added?
    is_removed: bool  # Has this VG been removed?
    res_changed: tuple[str, str] | None = None  # old vs new result type
    vg: VG  # Actual VG info


class ChangeSummary(BaseModel):
    """
    Summary of the formalization state changes since the last state where we had an admissable model
    """

    vg_changes: list[VGChange]

    def stats(self) -> dict[str, int]:
        """
        Return VG change statistics
        """

        def status_to_int(s: str) -> int:
            if s == 'unknown':
                return 1
            elif s == 'error' or s == 'err':
                return 2
            elif s == 'refuted':
                return 3
            elif s == 'verified_upto':
                return 4
            elif s == 'proved':
                return 5
            else:
                raise Exception(f'Unknown VG status: >{s}<')

        improved, deteriorated, proved, refuted = 0, 0, 0, 0
        for vg_change in self.vg_changes:
            if vg_change.is_removed or vg_change.res_changed is None:
                continue

            if vg_change.res_changed[1] in ['proved', 'verified_upto']:
                proved += 1
            if vg_change.res_changed[1] in ['refuted']:
                refuted += 1

            if not vg_change.is_new:
                if status_to_int(vg_change.res_changed[0]) > status_to_int(
                    vg_change.res_changed[1]
                ):
                    deteriorated += 1
                elif status_to_int(vg_change.res_changed[0]) > status_to_int(
                    vg_change.res_changed[1]
                ):
                    improved += 1

        return {
            'added': sum(1 if vgc.is_new else 0 for vgc in self.vg_changes),
            'removed': sum(1 if vgc.is_removed else 0 for vgc in self.vg_changes),
            'improved': improved,
            'deteriorated': deteriorated,
            'proved': proved,
            'refuted': refuted,
        }


def calc_vg_difference(old_vgs: list[VG], new_vgs: list[VG]) -> list[VGChange]:
    """
    Compute difference between two lists of VGs
    """
    vg_diff: list[VGChange] = []

    def find_a_match(new_v: VG, other_vgs: list[VG]) -> VG | None:
        def aux(new_data):
            def eq_preds(data):
                return data.predicate == new_data.predicate

            return find(lambda x: maybe_else(False, eq_preds, x.data), other_vgs)

        return maybe(aux, new_v.data)

    for new_vg in new_vgs:
        if new_vg.res is None:
            continue

        existing_vg = find_a_match(new_vg, old_vgs)

        if existing_vg is None:
            vg_diff.append(
                VGChange(
                    is_new=True,
                    is_removed=False,
                    res_changed=('', new_vg.res.res_type),
                    vg=new_vg,
                )
            )
        else:
            # This is an existing VG, let's see if its result type has changed at all...
            if (
                existing_vg.res is not None
                and new_vg.res.res_type != existing_vg.res.res_type
            ):
                vg_diff.append(
                    VGChange(
                        is_new=False,
                        is_removed=False,
                        res_changed=(existing_vg.res.res_type, new_vg.res.res_type),
                        vg=new_vg,
                    )
                )
            # ok, nothing changed
            else:
                vg_diff.append(
                    VGChange(
                        is_new=False, is_removed=False, res_changed=None, vg=new_vg
                    )
                )

    for old_vg in old_vgs:
        if old_vg.res is None:
            continue

        if find_a_match(old_vg, new_vgs) is None:
            vg_diff.append(
                VGChange(
                    is_new=False,
                    is_removed=True,
                    res_changed=(old_vg.res.res_type, ''),
                    vg=old_vg,
                )
            )

    return vg_diff


class SketchState(BaseModel):
    """
    SketchState contains the current version of
    """

    state_id: int  # State ID
    change: SketchChange  # Change that caused this state
    iml_code: str  # Actual IML code
    status: FormalizationStatus  # Formalization status
    error: str | None = None  # Error feedback from ImandraX

    vgs: list[VG]  # List of current verification goals
    old_vgs: list[VG]  # List of VGs in the last (good) state with admissable model
    decomps: list[RegionDecomp] = []  # Dictionary of decomposition requests/results

    vg_changes: list[VGChange] = Field(
        default_factory=lambda data: calc_vg_difference(
            old_vgs=data['old_vgs'],
            new_vgs=data['vgs'],
        )
        if 'vgs' in data and 'old_vgs' in data
        else []
    )

    # Summary of VG changes
    summary: ChangeSummary = Field(
        default_factory=lambda data: ChangeSummary(vg_changes=data['vg_changes'])
    )


class Sketch(BaseModel):
    """
    Sketch represent versions of a model
    """

    sketch_id: str = Field(default_factory=lambda: str(uuid.uuid4())[:8])
    processing: bool = False  # TODO do we need this for the UI?
    anchor_model_path: str  # relative path of the model that was used as the anchor
    init_iml_model: str = ''  # the initial model
    sketch_states: list[SketchState] = []  # here we keep track of changes in the state
    base_models: dict[
        str, Model
    ]  # Dict (by relative path) of base models that went into this
    def_owner: dict[str, str] = Field(
        default_factory=dict
    )  # def_name -> rel_path (e.g. "a.py")

    def get_state_by_idx(self, idx: int) -> SketchState:
        """
        Return SketchState for specific state ID
        """

        if not (0 <= idx and idx < len(self.sketch_states)):
            raise ValueError(f'Out of bounds state id: {idx}')

        return self.sketch_states[idx]

    def rollback(self, target_state_id: int | None = None):
        """
        rollback states
        """

        if target_state_id is None:
            # We should just be able to
            if len(self.sketch_states) in [0, 1]:
                raise Exception(
                    f'Not enough states[={len(self.sketch_states)}] to rollback!'
                )
            self.sketch_states = self.sketch_states[1:]

        else:
            if not (0 <= target_state_id and target_state_id < len(self.sketch_states)):
                raise Exception(
                    f'Invalid target state id[={target_state_id}] specified!'
                )

            self.sketch_states = self.sketch_states[: (target_state_id + 1)]

    def state_ids(self) -> list[int]:
        """
        Return list of SketchState IDs
        """
        if self.sketch_states:
            return list(range(len(self.sketch_states)))
        return []

    def latest_state(self) -> SketchState | None:
        """
        Return the latest sketch state if it's available
        """
        return self.sketch_states[0] if len(self.sketch_states) else None

    def last_good_state(self) -> SketchState | None:
        """
        Return last sketch state that has an admissable model
        """
        for i in range(len(self.sketch_states) - 1, -1, -1):
            if self.sketch_states[i].status == FormalizationStatus.TRANSPARENT:
                return self.sketch_states[i]

        return None

    def gen_init_task(self) -> SketchChangeTask:
        """
        This is the task that needs to be run to populate the initial state
        """
        return SketchChangeTask(
            sketch_id=self.sketch_id,
            iml_code=self.init_iml_model,
            change=SketchChgSetModel(new_iml_code=self.init_iml_model),
        )

    def get_task(self, change: SketchChange) -> SketchChangeTask:
        """
        Return a SketchChangeTask that can be sent to the worker to be analyzed by ImandraX
        """

        # Writing it this way so it's easier for pyright to pick up
        state = self.latest_state()

        if state is None:
            raise Exception('Attempting to create a task but we have no initial state!')

        new_iml_model = self.process_change(state.iml_code, change)

        # Now we have a sketch change task!
        return SketchChangeTask(
            sketch_id=self.sketch_id, change=change, iml_code=new_iml_model
        )

    def apply_change_result(self, change_result: SketchChangeResult) -> ChangeSummary:
        """
        Apply the result of running change task to the state and get the ChangeSummary as the result
        """
        if len(self.sketch_states) == 0 or self.last_good_state() is None:
            state_id = 0
            old_vgs = []
            self.init_iml_model = change_result.task.iml_code
        else:
            state_id = len(self.sketch_states) + 1
            old_vgs = self.last_good_state().vgs  # pyright: ignore

        new_state = SketchState(
            state_id=state_id,
            change=change_result.task.change,
            iml_code=change_result.task.iml_code,
            status=FormalizationStatus.TRANSPARENT
            if change_result.success
            else FormalizationStatus.INADMISSIBLE,
            error=change_result.error,
            vgs=change_result.vgs,
            old_vgs=old_vgs,
            decomps=change_result.decomps,
        )

        self.sketch_states.insert(0, new_state)

        return new_state.summary

    def _render_defs(self, defs: dict[str, str], order: list[str]) -> str:
        if not order:
            return ''
        # bodies already tend to end with "\n" from extract_definitions
        text = ''.join(defs[n] for n in order).rstrip() + '\n'
        return text

    def get_model_changes(self, *, misc_file_name: str = 'misc.iml') -> dict[str, str]:
        st = self.last_good_state()
        if st is None:
            return {}

        combined = st.iml_code.strip()
        combined_defs = extract_definitions(combined)
        combined_defs = {name: _norm_def(body) for name, body in combined_defs.items()}

        # Originals per-file, extracted from base models
        originals: dict[str, str] = {}
        orig_defs_by_file: dict[str, dict[str, str]] = {}

        for rel_path, m in self.base_models.items():
            p = _iml_path(rel_path)
            text = m.iml_model() or m.iml_code() or ''
            originals[p] = _ensure_trailing_nl(text)
            orig_defs_by_file[p] = extract_definitions(text)

        # Track per-file updated defs
        updated_defs_by_file = {
            p: dict(orig_defs_by_file[p]) for p in orig_defs_by_file
        }

        # Apply modifications back to owner file
        for name, new_body in combined_defs.items():
            owner_rel = self.def_owner.get(name)
            if owner_rel is None:
                # new definition: goes to misc
                continue

            owner_file = _iml_path(owner_rel)
            old_body = orig_defs_by_file.get(owner_file, {}).get(name)

            if old_body is None:
                # definition existed in combined but not found in owner file; treat as new -> misc
                continue

            if _norm_def(old_body) != new_body:
                updated_defs_by_file[owner_file][name] = new_body

        # Handle deletions: if a def existed in owner file but is not in combined, remove it
        for name, owner_rel in self.def_owner.items():
            owner_file = _iml_path(owner_rel)
            if (
                name in orig_defs_by_file.get(owner_file, {})
                and name not in combined_defs
            ):
                updated_defs_by_file[owner_file].pop(name, None)

        # New defs -> misc.iml
        new_defs = {n: b for n, b in combined_defs.items() if n not in self.def_owner}
        if new_defs:
            misc_text = originals.get(misc_file_name, '')
            misc_defs = extract_definitions(misc_text)
            misc_defs.update(new_defs)
            sorted_misc = topo_sort_definitions(misc_defs) if misc_defs else []
            originals[misc_file_name] = _ensure_trailing_nl(
                ''.join(misc_defs[n] for n in sorted_misc)
            )

        # Re-render each file in topo order and emit only diffs
        out: dict[str, str] = {}
        for file_path, defs in updated_defs_by_file.items():
            sorted_names = topo_sort_definitions(defs) if defs else []
            new_text = self._render_defs(defs, sorted_names)
            old_text = originals.get(file_path, '')
            if new_text != old_text:
                out[file_path] = new_text

        # If misc got new defs and differs
        if misc_file_name in originals and misc_file_name not in out:
            # if we created/updated misc above, compare to old (need old snapshot; simplest: only write if new_defs)
            if new_defs:
                out[misc_file_name] = originals[misc_file_name]

        return out

    def process_change(self, iml_code: str, change: SketchChange) -> str:
        """
        Process the change command for the specified code and return its updated version
        """

        if isinstance(change, SketchChgSetModel):
            updated_iml_code = change.new_iml_code

        elif isinstance(change, SketchChgInsertDef):
            try:
                updated_iml_code = add_definition(
                    source=iml_code, new_def=change.new_def_code
                )
            except Exception as e:
                raise Exception(f'Failed to insert: {e}')

        elif isinstance(change, SketchChgModifyDef):
            try:
                if isinstance(change.def_name_or_vg_idx, int):
                    updated_iml_code = update_definition(
                        source=iml_code,
                        name='',
                        new_def=change.new_def_body,
                        verify_index=change.def_name_or_vg_idx,
                    )
                else:
                    updated_iml_code = update_definition(
                        source=iml_code,
                        name=change.def_name_or_vg_idx,
                        new_def=change.new_def_body,
                    )
            except Exception as e:
                raise Exception(f'Failed to update definition: {e}')

        elif isinstance(change, SketchChgDeleteDef):
            try:
                if isinstance(change.def_name_or_vg_idx, int):
                    updated_iml_code = remove_definition(
                        source=iml_code, name='', verify_index=change.def_name_or_vg_idx
                    )
                else:
                    updated_iml_code = remove_definition(
                        source=iml_code, name=change.def_name_or_vg_idx
                    )
            except Exception as e:
                raise Exception(f'Failed to delete definition: {e}')

        else:
            raise Exception(f'Unsupported change event: {change}')

        return updated_iml_code

    def _baseline_iml_by_file(self) -> dict[str, str]:
        """
        Build the 'originals' dict for sync_definitions_topo.

        Keys here should be stable file identifiers. I recommend using `.iml` paths
        derived from each base model's rel_path.
        """
        originals: dict[str, str] = {}
        for rel_path, m in self.base_models.items():
            # Prefer "model-only" IML if you mean to patch defs, not artifacts.
            # If you only have iml_code(), that's ok too.
            src = m.iml_model() or m.iml_code() or ''
            iml_path = str(Path(rel_path).with_suffix('.iml'))
            originals[iml_path] = src.strip()
        return originals

    def __repr__(self):
        return f'Sketch [ID = {self.sketch_id}; Anchor model path = {self.anchor_model_path}]'

    @staticmethod
    def build_from_models(models: list[Model], *, anchor_model_path: str) -> 'Sketch':
        base_models = {m.rel_path: m for m in models}

        combined_defs: dict[str, str] = {}
        def_owner: dict[str, str] = {}
        dup_owners: dict[str, set[str]] = {}

        for m in models:
            src = (m.iml_model() or m.iml_code() or '').strip()
            if not src:
                continue

            defs = extract_definitions(src)
            for name, body in defs.items():
                # Detect ambiguous ownership (same top-level name defined in >1 base model)
                existing_owner = def_owner.get(name)
                if existing_owner is not None and existing_owner != m.rel_path:
                    dup_owners.setdefault(name, set()).update(
                        {existing_owner, m.rel_path}
                    )
                    continue

                # Keep first body we saw (content doesn't matter once ownership is unique)
                if name not in combined_defs:
                    combined_defs[name] = body
                    def_owner[name] = m.rel_path

        if dup_owners:
            parts: list[str] = []
            for name in sorted(dup_owners.keys()):
                owners = ', '.join(sorted(dup_owners[name]))
                parts.append(f'{name}: {owners}')
            raise ValueError(
                'Duplicate top-level definition names across base models: '
                + '; '.join(parts)
            )

        sorted_names = topo_sort_definitions(combined_defs)
        init_iml_model = ''.join(combined_defs[n] for n in sorted_names).strip() + '\n'

        sk = Sketch(
            anchor_model_path=anchor_model_path,
            init_iml_model=init_iml_model,
            base_models=base_models,
        )
        sk.def_owner = def_owner
        return sk
