"""The Matrix Engine: Central Intelligence Layer.

Detects active modules, builds a dependency DAG of integrations, computes
the OptimizationProfile, and generates the consolidated integration layer.
"""

from __future__ import annotations

import importlib
import pkgutil
from collections import defaultdict, deque
from pathlib import Path
from typing import TYPE_CHECKING, Any

from nikame.codegen.integrations.base import BaseIntegration, OptimizationProfile
from nikame.utils.logger import get_logger

if TYPE_CHECKING:
    from nikame.blueprint.engine import Blueprint
    from nikame.config.schema import NikameConfig
    from nikame.utils.file_writer import FileWriter

_log = get_logger("matrix_engine")


class MatrixEngine:
    """Orchestrates all cross-module capabilities dynamically."""

    def __init__(self, config: NikameConfig, blueprint: Blueprint, writer: FileWriter) -> None:
        self.config = config
        self.blueprint = blueprint
        self.writer = writer
        
        self.active_modules = {m.NAME for m in blueprint.modules}
        self.active_features = set(config.features)
        
        # Integration class registry: Name -> Class
        self._registry: dict[str, type[BaseIntegration]] = {}
        
        self._discover_integrations()

    def _discover_integrations(self) -> None:
        """Auto-discover all BaseIntegration and BaseCodegen subclasses."""
        # 1. Discover BaseIntegration subclasses in the integrations package
        package_path = Path(__file__).parent.resolve()
        package_name = "nikame.codegen.integrations"

        for _, modname, _ in pkgutil.walk_packages([str(package_path)], prefix=f"{package_name}."):
            if modname.split(".")[-1] in ("base", "matrix"):
                continue

            try:
                module = importlib.import_module(modname)
                for attr_name in dir(module):
                    attr = getattr(module, attr_name)
                    if (
                        isinstance(attr, type)
                        and issubclass(attr, BaseIntegration)
                        and attr is not BaseIntegration
                    ):
                        self._registry[attr.__name__] = attr
            except Exception as e:
                _log.warning(f"Failed to import integration module {modname}: {e}")
        
        # 2. Add BaseCodegen subclasses from COMPONENT_REGISTRY
        from nikame.codegen.registry import COMPONENT_REGISTRY
        for key, info in COMPONENT_REGISTRY.items():
            cls = info.get("class")
            if cls:
                self._registry[cls.__name__] = cls

    def _compute_profile(self) -> OptimizationProfile:
        """Compute system variables based on project metadata."""
        scale = self.config.project.scale
        pattern = self.config.project.access_pattern
        
        # Scale sizing
        if scale == "small":
            max_conn, workers, partitions = 10, 2, 3
        elif scale == "large":
            max_conn, workers, partitions = 100, 16, 12
        else: # medium
            max_conn, workers, partitions = 25, 4, 6
            
        # Detect capabilities from features/modules
        is_rag = any(f in self.active_features for f in ["rag-pipeline", "semantic-search"])
        is_pipeline = any(f in self.active_features for f in ["data-pipeline", "streaming-ingestion"])
        is_high_perf = any(m in self.active_modules for m in ["dragonfly", "clickhouse", "redpanda"])

        if is_rag:
            # RAG apps are read-heavy on vectors, write-heavy on logs
            max_conn += 20 
        if is_pipeline:
            # Pipelines need more partitions
            partitions *= 2
            workers *= 2
        if is_high_perf:
            max_conn = int(max_conn * 1.5)

        # Pattern tuning
        if pattern == "read_heavy":
            cache_ttl = 3600 * 24 # 24hrs
        elif pattern == "write_heavy":
            cache_ttl = 300 # 5min
        else: # balanced
            cache_ttl = 3600 # 1hr

        return OptimizationProfile(
            project_scale=scale,
            access_pattern=pattern,
            max_connections=max_conn,
            cache_ttl_seconds=cache_ttl,
            kafka_partitions=partitions,
            worker_concurrency=workers
        )

    def _topological_sort(self, triggered: list[Any]) -> list[Any]:
        """Sort implementations dynamically via their dependencies."""
        name_to_inst = {inst.__class__.__name__: inst for inst in triggered}
        
        # Also map by string NAME for BaseCodegen dependencies
        str_name_to_cls_name = {}
        for inst in triggered:
            if hasattr(inst, "NAME"):
                str_name_to_cls_name[inst.NAME] = inst.__class__.__name__
        
        # Build graph
        graph: dict[str, list[str]] = defaultdict(list)
        in_degree: dict[str, int] = {inst.__class__.__name__: 0 for inst in triggered}
        
        for inst in triggered:
            cls_name = inst.__class__.__name__
            
            deps = getattr(inst, "DEPENDS_ON", [])
            str_deps = getattr(inst, "DEPENDENCIES", [])
            
            all_deps = list(deps)
            for d in str_deps:
                if d in str_name_to_cls_name:
                    all_deps.append(str_name_to_cls_name[d])

            for dep_cls in all_deps:
                if dep_cls in name_to_inst:
                    graph[dep_cls].append(cls_name)
                    in_degree[cls_name] += 1
                    
        # Kahn's algorithm
        queue = deque([k for k, v in in_degree.items() if v == 0])
        sorted_names = []
        
        while queue:
            node = queue.popleft()
            sorted_names.append(node)
            for neighbor in graph[node]:
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    queue.append(neighbor)
                    
        if len(sorted_names) != len(triggered):
            _log.error("Matrix engine detected a circular dependency in integrations!")
            # Fallback to an arbitrary order if circle is found (should be avoided)
            return triggered
            
        return [name_to_inst[name] for name in sorted_names]

    def execute(self) -> None:
        """Discover, score, sort, and dispatch integrations."""
        profile = self._compute_profile()
        _log.info(f"Matrix computed profile: Scale '{profile.project_scale}', Access '{profile.access_pattern}'")

        # 1. First, instantiate all triggered integrations
        triggered_instances = []
        for name, cls in self._registry.items():
            if cls.should_trigger(self.active_modules, self.active_features):
                # Instantiate with computed profile or appropriate signatures
                if issubclass(cls, BaseIntegration):
                    instance = cls(self.config, self.blueprint, profile)
                else:
                    # BaseCodegen signature: (ctx, config)
                    from nikame.codegen.base import CodegenContext
                    ctx = CodegenContext(
                        project_name=self.config.name,
                        active_modules=list(self.active_modules),
                        features=list(self.active_features)
                    )
                    instance = cls(ctx, self.config)
                
                triggered_instances.append(instance)
                _log.debug(f"Matrix flagged: {name}")

        if not triggered_instances:
            _log.debug("No matrix integrations triggered.")
            return

        # 2. Sort the INSTANCES based on their class DEPENDS_ON rules
        sorted_instances = self._topological_sort(triggered_instances)

        _log.info(f"Matrix dispatching {len(sorted_instances)} configured integrations.")
        
        # Aggregated outputs to be injected into main systems
        lifespans = []
        health_checks = {}
        metrics_blocks = []
        guides = []
        
        for integration in sorted_instances:
            # 1. Write core output files
            # Check if it's a BaseIntegration or BaseCodegen (they have different method names)
            if isinstance(integration, BaseIntegration):
                for path, content in integration.generate_core():
                    _log.info(f"Matrix writing integration file: {path}")
                    self.writer.write_file(path, content)
                
                # 2. Collect aggregation blocks
                lifespans.append(integration.generate_lifespan())
                health_checks.update(integration.generate_health())
                metrics_blocks.append(integration.generate_metrics())
                guides.append(integration.generate_guide())
            else:
                # BaseCodegen style
                for path, content in integration.generate():
                    _log.info(f"Matrix writing component file: {path}")
                    self.writer.write_file(path, content)
                
                # For components, guide metadata is handled slightly differently
                # but we can still extract it
                guides.append(f"### {integration.NAME}\n{integration.DESCRIPTION}\n")

        # 3. Inject the collected components into main templates
        self._inject_into_app(lifespans, health_checks, metrics_blocks)
        
        # 4. Integrate docs directly
        if guides:
            # We assume a mechanism to write the guide back to documentation layer
            # which could be handled natively or passing back to guide generator.
            combined_guide = "\n\n".join(guides)
            self.writer.write_file("docs/integrations.md", combined_guide)
            
    def _inject_into_app(
        self,
        lifespans: list[str],
        health_checks: dict[str, str],
        metrics: list[str]
    ) -> None:
        """Inject collected blocks into the main application file."""
        main_py = "services/api/app/main.py"
        # In a real system, we'd more robustly resolve the main.py path
        # For now, we use the standard NIKAME structure or check if it exists in writer
        
        # Collect existing content
        content = ""
        try:
            # Check if writer has the file in its buffer or on disk
            # FileWriter.read_file() doesn't exist, we must rely on what was written
            # Actually, MatrixEngine should probably work on the file buffer before it's flushed
            # but FileWriter writes immediately to disk.
            # So we read from disk.
            target_path = Path(self.writer.output_dir) / main_py
            if not target_path.exists():
                _log.warning(f"Matrix could not find {main_py} for injection")
                return
            content = target_path.read_text()
        except Exception as e:
            _log.error(f"Matrix failed to read {main_py}: {e}")
            return

        lines = content.splitlines()

        # 1. Inject Lifespan Hooks
        if lifespans:
            # Find the lifespan function if it exists
            # For simplicity, we assume the AutoWiringEngine might have created one,
            # or we look for a marker.
            insert_idx = -1
            for i, line in enumerate(lines):
                if "async def lifespan" in line:
                    insert_idx = i + 1
                    break
            
            if insert_idx != -1:
                # Find yield in lifespan
                yield_idx = -1
                for i in range(insert_idx, len(lines)):
                    if "yield" in lines[i]:
                        yield_idx = i
                        break
                
                if yield_idx != -1:
                    for hook in lifespans:
                        hook_line = f"    {hook}"
                        # Check if it already exists (prevent duplication on regenerate)
                        is_duplicate = False
                        for line in lines:
                            if hook in line:
                                is_duplicate = True
                                break
                        if not is_duplicate:
                            lines.insert(yield_idx, hook_line)
                            yield_idx += 1
            else:
                _log.warning("No lifespan found in main.py for Matrix injection")

        # 2. Inject Health Checks
        if health_checks:
            # Aggregator handled by AutoWiringEngine usually, but Matrix
            # can inject into the specific integration health map if needed.
            pass # Currently handled better by AutoWiringEngine Pass 3

        # 3. Inject Metrics
        if metrics:
            # Metrics usually go into a Prometheus route or middleware
            pass

        # Write back
        new_content = "\n".join(lines)
        if new_content != content:
            self.writer.write_file(main_py, new_content)
