#
#   Imandra Inc.
#
#   codelogician/modeling/metamodel.py
#

import logging
from collections import defaultdict
from pathlib import Path

from imandra.u.agents.code_logician.base import FormalizationStatus
from pydantic import BaseModel, Field, ValidationInfo, model_validator

from ..server.events import FileSystemEvent, FileSystemEventType
from ..sketch.sketch import Sketch
from ..strategy.cl_agent_state import CLAgentState, CLResult
from .dep_context import ModelDepContext
from .fs_sync import model_filesystem_sync
from .model import Model
from .task import ModelTask, TaskKind
from .task_creator import ModelTaskCreator

log = logging.getLogger(__name__)


class CycleSummary(BaseModel):
    rel_paths: list[str]
    edge_paths: list[tuple[str, str]]


class MetaModel(BaseModel):
    """
    Container for the set of models representing the whole project
    """

    src_dir_abs_path: str
    models: dict[str, Model]
    language: str = 'Python'

    dep_cycles: list[CycleSummary] = Field(default_factory=list)

    @model_validator(mode='after')
    def link_model_deps(self, info: ValidationInfo) -> 'MetaModel':
        # Only rebuild edges when explicitly asked (e.g., load-from-disk path).
        if not (info.context or {}).get('relink_from_str_edges', False):
            return self

        # Clear in-memory edges
        for m in self.models.values():
            m.dependencies.clear()
            m.rev_dependencies.clear()

        # Rebuild from persisted strings (deps first)
        for m in self.models.values():
            for dep_path in m.str_dependencies:
                dep = self.models.get(dep_path)
                if dep is not None:
                    m.dependencies.append(dep)

        # Rebuild rev-deps as inverse to guarantee consistency
        for m in self.models.values():
            for dep in m.dependencies:
                dep.rev_dependencies.append(m)

        return self

    def get_paths_with_dirs(self):
        """
        Return a list of paths with directory information
        """

        FILE_MARKER = '<files>'

        def attach(branch, trunk, model):
            parts = branch.split('/', 1)
            if len(parts) == 1:  # branch is a file
                trunk[FILE_MARKER].append(model)
            else:
                node, others = parts
                if node not in trunk:
                    trunk[node] = defaultdict(dict, ((FILE_MARKER, []),))
                attach(others, trunk[node], model)

        main_dict = defaultdict(dict, ((FILE_MARKER, []),))

        for path in sorted(self.models.keys()):
            attach(path, main_dict, self.models[path])

        return main_dict

    def list_models(self) -> list[Model]:
        """
        Return the list of all models
        """
        return list(self.models.values())

    def vgs(self):
        """
        Return all verification goals for all the models
        """
        return {m.rel_path: m.verification_goals() for m in self.models.values()}

    def decomps(self):
        """
        Return decomposition requests for all the models
        """
        return {m.rel_path: m.decomps() for m in self.models.values()}

    def opaques(self):
        """
        Return all the opaque functions for all the models
        """
        return {m.rel_path: m.opaque_funcs() for m in self.models.values()}

    def apply_cl_result(self, cl_result: CLResult) -> None:
        """
        Applies the result of running CodeLogician on a model
        """

        if cl_result.task.rel_path in self.models:
            model = self.models[cl_result.task.rel_path]

            # Rememeber that we may have multiple tasks outstanding simultaneously
            # so, we always keep the last one that we use
            if model.outstanding_task_ID == cl_result.task.task_id:
                dependencies = cl_result.task.dependencies
                model.apply_agent_state(cl_result.agent_state, dependencies)
                model.outstanding_task_ID = None
                model.outstanding_task_kind = None

                if cl_result.task.task_kind == TaskKind.ANALYZE:
                    model.analysis_vgs_attempted |= bool(cl_result.task.gen_vgs)
                    model.analysis_decomps_attempted |= bool(cl_result.task.gen_decomps)

                elif cl_result.task.task_kind == TaskKind.USER_IML_UPDATE:
                    # User edits implicitly “finalize” analysis for this baseline
                    model.analysis_vgs_attempted = True
                    model.analysis_decomps_attempted = True

                elif cl_result.task.task_kind == TaskKind.FORMALIZE:
                    # New baseline → allow analyze again
                    model.analysis_vgs_attempted = False
                    model.analysis_decomps_attempted = False

                log.info('CLResult applied to model successfully!')
            else:
                log.warning(
                    f'Model with path={cl_result.task.rel_path} has a newer Task ID [{model.outstanding_task_ID}], '
                    f' so ignoring this one [{cl_result.task.task_kind}:{cl_result.task.task_id}]!'
                )
        else:
            log.error(
                f'Attempting to apply CL Result to a non-existent model path: {cl_result.task.rel_path}!'
            )

    def update_model(self, rel_path: str, state: CLAgentState) -> None:
        """
        Update the model with a new agent state (formalization state form CL agent)
        """

        if rel_path in self.models:
            self.models[rel_path].apply_agent_state(state)
            log.info(
                f'AgentState successfully updated for model with path=[{rel_path}]'
            )
        else:
            log.error(
                f'Attempting to set agent state to a non-existent model path {rel_path}'
            )

    def get_next_tasks(self, user_wait_time: int | None = None) -> list[ModelTask]:
        log.debug('About to get_next_tasks')

        # Hard gate: do not schedule tasks if cycles exist
        if getattr(self, 'dep_cycles', []):
            log.warning(
                'Not scheduling tasks because dependency cycles are present: %s',
                [c.rel_paths for c in self.dep_cycles],
            )
            return []

        tasks: list[ModelTask] = []

        task_creator = ModelTaskCreator(user_wait_time=user_wait_time)

        for model in self.models.values():
            t = task_creator.create_task(model)
            if t is not None:
                # Record outstanding task info for stale-result detection
                model.outstanding_task_ID = t.task_id
                model.outstanding_task_kind = t.task_kind
                tasks.append(t)

                log.debug(
                    'Scheduled task: path=%s task_id=%s task_kind=%s class=%s '
                    '(gen_vgs=%s gen_decomps=%s)',
                    model.rel_path,
                    t.task_id,
                    t.task_kind.name if t.task_kind else None,
                    type(t).__name__,
                    getattr(t, 'gen_vgs', None),
                    getattr(t, 'gen_decomps', None),
                )

        if tasks:
            log.debug(
                'get_next_tasks summary: %s',
                [
                    {
                        'path': t.rel_path,
                        'task_kind': t.task_kind.name if t.task_kind else None,
                        'class': type(t).__name__,
                        'task_id': t.task_id,
                    }
                    for t in tasks
                ],
            )
        else:
            log.debug('get_next_tasks summary: no tasks scheduled')

        log.debug('Returning %d tasks', len(tasks))
        return tasks

    def abs_path_in_models(self, path: str) -> bool:
        """
        Return True if the specified path in the models
        """
        return self.make_path_relative(path) in self.models

    def make_path_relative(self, path: str) -> str:
        """
        Return True/False if the specified path is relative to the base source code directory
        """
        return str(Path(path).relative_to(self.src_dir_abs_path))

    def get_model_by_path(self, rel_path: str) -> Model | None:
        """
        Returns Model object for specified relative path, None if it's missing
        """

        if rel_path in self.models:
            return self.models[rel_path]
        return None

    def gather_deep_dependencies(self, rel_path: str) -> list[Model]:
        """
        Return list of all "deep" dependencies for the specified relative path
        """

        if rel_path not in self.models:
            return []

        def rec_helper(curr_model: Model):
            res = []
            for d in curr_model.dependencies:
                assert isinstance(d, Model)
                res.extend(rec_helper(d) + [d])
            return res

        curr_model = self.models[rel_path]
        return list(set(rec_helper(curr_model) + [curr_model]))

    def gen_consolidated_model(self, rel_path: str) -> tuple[str, dict[str, Model]]:
        """
        Generate a single consolidated model from the achnor model specified by (rel_path).
        """

        if rel_path not in self.models:
            raise Exception(f'Specified model {rel_path} is not a model!')

        # we'll gather the relevant models (subgraph) and then sort it
        relevant_models: list[Model] = self.gather_deep_dependencies(rel_path)
        sorted_list: list[Model] = ModelDepContext.topological_sort(relevant_models)

        # now we'll consolidate the code
        iml = ''
        for model in sorted_list:
            iml += (
                f'(* Starting {model.rel_path} *)\n'
                + (str(model.iml_code()) if model.iml_code() else '')
                + f'\n(* Ending {model.rel_path} *)\n'
            )

        return iml, {m.rel_path: m for m in relevant_models}

    def process_filesystem_event(self, event: FileSystemEvent):
        """
        Note that the provided paths are absolute!
        We receive events that concern both - Python and IML files.
        """

        # Let's figure out if we have a source code file or we have
        ext = Path(event.abs_path1).suffix
        if ext == '.py':
            self.process_fs_event_src_code(event)
        elif ext == '.iml':
            self.process_fs_event_iml_code(event)
        else:
            log.warning(
                f'MetaModel asked to process FS event for unknown extension: {ext}'
            )

    def process_fs_event_iml_code(self, event: FileSystemEvent):
        """
        Process FileSystem event when it's for a model
        """

        if event.action_type in [
            FileSystemEventType.CREATED,
            FileSystemEventType.MODIFIED,
        ]:
            # Let's first figure out what the Model path would be for this guy

            rel_path = self.make_path_relative(event.abs_path1)
            model_path = str(Path(rel_path).with_suffix('.py'))

            if model_path in self.models:
                try:
                    with open(event.abs_path1) as infile:
                        imlCode = infile.read()
                except Exception as e:
                    log.error(f'Failed to read IML file: {event.abs_path1}: {e}')
                    return

                self.models[model_path].set_iml_model(imlCode, record_time=True)
                self.models[model_path].expected_iml_on_disk = imlCode

                log.info(f'Updated IML code for {model_path}')
            else:
                log.info(
                    f'Received IML file for unknown model. Implied model path is: {model_path}'
                )
        else:
            log.info(
                f'We only support CREATED/MODIFIED FS events for IML models. Got: [{event.action_type}]'
            )

    def process_fs_event_src_code(self, event: FileSystemEvent):
        """
        Process FileSystem event when it's for a source code file
        """

        if event.action_type == FileSystemEventType.CREATED:
            if self.make_path_relative(event.abs_path1) in self.models:
                log.warning(
                    'Attempting to create a model with a path that already exists!'
                )
                return
            else:
                log.info(f'Creating a model for {event.abs_path1}')

                abs_path = event.abs_path1
                rel_path = self.make_path_relative(abs_path)
                try:
                    with open(abs_path) as infile:
                        srcCode = infile.read()
                except Exception:
                    log.error(f'Failed to read contents of the created file {abs_path}')
                    return

                self.models[rel_path] = Model(rel_path=rel_path, src_code=srcCode)
                log.info(f"We've now created a new model for path={rel_path}.")

        elif event.action_type == FileSystemEventType.MODIFIED:
            # A source file has been modified, so we will update our model info
            log.info(f'Processing modified source file: {event.abs_path1}')

            abs_path = event.abs_path1
            rel_path = self.make_path_relative(abs_path)

            try:
                with open(abs_path) as infile:
                    srcCode = infile.read()
            except Exception as e:
                log.error(
                    f'Failed to open a modified file: fullpath={abs_path}; error={str(e)}'
                )
                return

            if rel_path in self.models:
                # Let's now update the model's source code and record when we've done it
                self.models[rel_path].set_src_code(srcCode, True)
                log.info(
                    f"We've now updated source code for model with path=[{abs_path}]"
                )
            else:
                log.info(
                    'Modified file does not have a corresponding model. Creating it.'
                )
                self.models[rel_path] = Model(rel_path=rel_path, src_code=srcCode)
                log.info(f"We've now created a new model for path={abs_path}.")

        elif event.action_type == FileSystemEventType.DELETED:
            # We need to process a deleted file event
            log.info(f'Processing deleted source file: {event.abs_path1}')
            rel_path = self.make_path_relative(event.abs_path1)
            if rel_path in self.models:
                self.models[rel_path].set_src_code(None)
            else:
                log.error(
                    f'Attempting to process delete event for a non-existent model {event.abs_path1}'
                )

        elif event.action_type == FileSystemEventType.MOVED:
            # We need to process a deleted file event
            log.info(f'Processing a moved file: {event.abs_path1} to {event.abs_path2}')
            rel_path1 = self.make_path_relative(event.abs_path1)
            rel_path2 = self.make_path_relative(event.abs_path2)

            if rel_path1 in self.models and rel_path2 not in self.models:
                # self._models[rel_path].setSrcCode(None)
                self.models[rel_path2] = self.models[rel_path1]
                self.models[rel_path2].rel_path = rel_path2
                del self.models[rel_path1]
        else:
            log.error(
                "Something went wrong. We've gotten an unknown event type: {event.action_type}"
            )

    @staticmethod
    def _canonicalize_cycle_paths(rel_paths: list[str]) -> list[str]:
        # rel_paths is a closed walk: p[0] == p[-1]
        if len(rel_paths) <= 2:
            return rel_paths

        core = rel_paths[:-1]  # drop closing repeat
        # pick rotation that starts with smallest rel_path
        min_i = min(range(len(core)), key=lambda i: core[i])
        rotated = core[min_i:] + core[:min_i]
        return rotated + [rotated[0]]

    @staticmethod
    def _edge_paths_from_rel_paths(rel_paths: list[str]) -> list[tuple[str, str]]:
        return list(zip(rel_paths, rel_paths[1:]))

    def recompute_dep_cycles(self) -> None:
        """
        Recompute dependency cycles based on current live model graph.
        """
        models = sorted(self.models.values(), key=lambda m: m.rel_path)
        cycles = ModelDepContext.detect_cycles_in_models(models)

        summaries: list[CycleSummary] = []
        for c in cycles:
            rp = self._canonicalize_cycle_paths(c.rel_paths)
            ep = self._edge_paths_from_rel_paths(rp)
            summaries.append(CycleSummary(rel_paths=rp, edge_paths=ep))

        summaries.sort(key=lambda s: (len(s.rel_paths), s.rel_paths))
        dedup: dict[tuple[str, ...], CycleSummary] = {}
        for s in summaries:
            dedup[tuple(s.rel_paths)] = s

        self.dep_cycles = list(dedup.values())

    def run_file_sync(self):
        """
        Go through the source code directory and update the models.
        """

        log.info(f'Building FS from {self.src_dir_abs_path}')
        try:
            self.models = model_filesystem_sync(
                models=self.models,
                src_dir_abs_path=self.src_dir_abs_path,
                language=self.language,
            )
        except Exception as e:
            msg = f'Failed to perform model_filesync: {e}'
            log.error(msg)
            raise Exception(msg)

        self.recompute_dep_cycles()

        if not self.dep_cycles:
            log.info('FileSync ran successfully! No cycles detected')
        else:
            log.info(f'FileSync finished. Cycles detected: {self.dep_cycles}')

    def apply_sketch(self, sketch: Sketch) -> None:
        """
        Apply sketch back to the MetaModel
        """

        raise NotImplementedError('Sketch application is not implemented yet!')

    @staticmethod
    def fromJSON(j: str | dict):
        """
        Return an object from JSON file
        """
        if isinstance(j, str):
            return MetaModel.model_validate_json(
                j, context={'relink_from_str_edges': True}
            )
        else:
            return MetaModel.model_validate(j, context={'relink_from_str_edges': True})

    def toJSON(self):
        """
        Return a JSON object
        """
        return self.model_dump_json()

    def summary(self):
        """Return summary of the meta model state"""
        allFormStatuses = [
            FormalizationStatus.UNKNOWN,
            FormalizationStatus.TRANSPARENT,
            FormalizationStatus.INADMISSIBLE,
            FormalizationStatus.EXECUTABLE_WITH_APPROXIMATION,
            FormalizationStatus.ADMITTED_WITH_OPAQUENESS,
        ]

        frmStatuses = {s: 0 for s in allFormStatuses}
        for model in self.models.values():
            frmStatuses[model.formalization_status()] += 1

        def count(f):
            return sum(len(f(m)) for m in self.models.values())

        return {
            'frm_statuses': frmStatuses,
            **{
                label: count(f)
                for label, f in [
                    ('num_models', lambda _: [None]),
                    ('num_vgs', lambda m: m.verification_goals()),
                    ('num_opaque_funcs', lambda m: m.opaque_funcs()),
                    ('num_decomps', lambda m: m.decomps()),
                ]
            },
        }

    def eq_key(self) -> dict:
        return {
            'src_dir_abs_path': self.src_dir_abs_path,
            'language': self.language,
            'dep_cycles': [c.model_dump() for c in self.dep_cycles],
            'models': {
                rel: {
                    'rel_path': m.rel_path,
                    'src_code': m.src_code,
                    'src_language': m.src_language,
                    'context': m.context,
                    'iml_code_frozen': m.iml_code_frozen,
                    'outstanding_task_ID': m.outstanding_task_ID,
                    'outstanding_task_kind': (
                        m.outstanding_task_kind.name
                        if m.outstanding_task_kind
                        else None
                    ),
                    'analysis_vgs_attempted': m.analysis_vgs_attempted,
                    'analysis_decomps_attempted': m.analysis_decomps_attempted,
                    'str_dependencies': sorted(list(m.str_dependencies)),
                }
                for rel, m in sorted(self.models.items(), key=lambda kv: kv[0])
            },
        }

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, MetaModel):
            return NotImplemented
        return self.eq_key() == other.eq_key()
