src/runtime.py

"""
MORPH-AI Runtime v6
Local inference engine with skill management.
Supports CPU, GPU, Apple Silicon, mobile (via GGUF/Termux), AND v6 capabilities:
- Audio: ASR (Whisper) + TTS (Coqui/gTTS)
- Tools: JSON-structured function calling
- Documents: PDF/DOCX/OCR extraction
- Video: Frame sampling + temporal reasoning
- Multimodal: Vision + Audio + Video fusion
"""

import json
import base64
import os
import sys
import torch
import torch.nn.functional as F
from pathlib import Path
from typing import Optional, List, Dict, Union
from dataclasses import dataclass
import numpy as np

_ROOT = Path(__file__).resolve().parent.parent
if str(_ROOT / "src") not in sys.path:
    sys.path.insert(0, str(_ROOT / "src"))

try:
    from llama_cpp import Llama
    HAS_LLAMACPP = True
except Exception:
    Llama = None
    HAS_LLAMACPP = False

from architecture import MorphModel, MorphConfig, build_code_features
from fsm import RuntimeFSM, IllegalTransitionError
from rules import RuleEngine
from regex_features import RegexFeatureExtractor
from routing import RoutingMatrix, deterministic_index
from kvstore import KVStore
from vision import VisionAnalyzer, ImageFacts
from search import SearchClient, RAGPipeline
from facts import FactExtractor, KnowledgeGraph, GraphQuery
from audio import AudioModule, AudioFacts
from document import DocumentModule
from video import VideoModule
from tools import ToolRegistry, ToolCall


@dataclass
class Skill:
    name: str
    token: str
    description: str
    trigger_patterns: List[str]
    lora_weights: Optional[bytes] = None
    examples: List[Dict[str, str]] = None
    installed: bool = False

    def __post_init__(self):
        if self.examples is None:
            self.examples = []

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "token": self.token,
            "description": self.description,
            "trigger_patterns": self.trigger_patterns,
            "lora_weights": base64.b64encode(self.lora_weights).decode() if self.lora_weights else None,
            "examples": self.examples,
            "installed": self.installed
        }

    @classmethod
    def from_dict(cls, data: dict) -> 'Skill':
        lora_bytes = base64.b64decode(data["lora_weights"]) if data.get("lora_weights") else None
        return cls(
            name=data["name"],
            token=data["token"],
            description=data["description"],
            trigger_patterns=data["trigger_patterns"],
            lora_weights=lora_bytes,
            examples=data.get("examples", []),
            installed=data.get("installed", False)
        )

    def save(self, path: str):
        with open(path, 'w') as f:
            json.dump(self.to_dict(), f, indent=2)

    @classmethod
    def load(cls, path: str) -> 'Skill':
        with open(path, 'r') as f:
            data = json.load(f)
        return cls.from_dict(data)


class MorphRuntime:
    """
    Runtime engine for MORPH-AI v6.
    Handles model loading, inference, skill management, and v6 capabilities:
    Audio (ASR + TTS), Tools, Documents, Video, Multimodal fusion.
    """

    def __init__(self, model_path: str, use_4bit: bool = True, use_cpu: bool = False, plugin_dir: Optional[str] = None):
        self.model_path = Path(model_path)
        self.use_4bit = use_4bit and not use_cpu
        self.use_cpu = use_cpu
        self.skills: Dict[str, Skill] = {}
        self.active_skill: Optional[str] = None
        self.plugin_dir = plugin_dir

        # v5 pipeline layers
        self.fsm = RuntimeFSM()
        self.rules = RuleEngine(str(_ROOT / "rules" / "rules.json"))
        self.regex = RegexFeatureExtractor(num_features=7)
        self.routing = RoutingMatrix(str(_ROOT / "routing" / "routing_matrix.json"))
        self.kv = KVStore(str(_ROOT / "cache" / "kvstore.json"))

        # v5.5 multimodal + live knowledge layers
        self.vision = VisionAnalyzer()
        self.search_client = SearchClient()
        self.rag = RAGPipeline(client=self.search_client, cache=self.kv)
        self.fact_extractor = FactExtractor()
        self.graph = KnowledgeGraph()
        self.gquery = GraphQuery(self.fact_extractor)

        # v6 advanced capabilities
        self.audio = AudioModule(MorphConfig(), self._get_hidden_dim())
        self.document = DocumentModule(MorphConfig(), self._get_hidden_dim())
        self.video = VideoModule(MorphConfig(), self._get_hidden_dim())
        self.tool_registry = ToolRegistry()

        print(f"Loading MORPH-AI model from {model_path}")
        self._load_model()

    def _get_hidden_dim(self) -> int:
        try:
            return self.model.base_model_raw.config.hidden_size
        except Exception:
            return 768

    def _transition(self, next_state: str):
        try:
            self.fsm.transition(next_state)
        except IllegalTransitionError as e:
            self.fsm.fault()
            raise

    # ------------------------------------------------------------------
    # v6 multimodal ingest path
    # GUARD_IN -> VISION (if image) -> AUDIO (if audio) -> VIDEO (if video) ->
    # DOCUMENT (if document) -> SEARCH_GATE -> SEARCH -> FACT_EXTRACT
    # ------------------------------------------------------------------
    _LIVE_NEED_HINTS = ("today", "latest", "current", "news", "price", "weather",
                        "live", "breaking", "who won", "election", "stock", "score")

    def _needs_live_data(self, prompt: str, force_search: bool, depth: int) -> bool:
        if force_search:
            return True
        if depth >= 2:
            return False
        low = prompt.lower()
        return any(h in low for h in self._LIVE_NEED_HINTS)

    def _ingest(self, prompt: str, image_path: Optional[str] = None,
                audio_path: Optional[str] = None, video_path: Optional[str] = None,
                document_path: Optional[str] = None,
                force_search: bool = False, depth: int = 0) -> dict:
        """Run guard-in + (optional) vision + audio + video + document + search/facts.

        Returns a context dict consumed by chat / chat_best_of_n:
          image_facts, audio_facts, video_facts, document_text, search_context, graph_context, enriched.
        """
        # L1 in-bound guardrails
        in_dec = self.rules.eval(prompt, phase="in")
        if in_dec.action == "block":
            self._transition("RESPOND")
            self._transition("IDLE")
            return {"blocked": in_dec.reply or "I can't help with that."}

        syn = self.regex.extract_text(prompt)
        gate = self.regex.gate(syn)
        if gate == "block":
            self._transition("RESPOND")
            self._transition("IDLE")
            return {"blocked": "That input couldn't be processed (unbalanced syntax)."}
        if gate == "warn":
            print(f"warn: syntax state {syn.to_dict()}")

        image_facts: Optional[ImageFacts] = None
        image_text = ""
        audio_facts = None
        audio_text = ""
        video_facts = None
        video_text = ""
        document_text = ""

        # v6: multimodal phases
        if image_path:
            self._transition("VISION")
            try:
                image_facts = self.vision.analyze(image_path)
                image_text = image_facts.to_text()
            except Exception as e:
                print(f"vision failed (non-fatal): {e}")

        if audio_path:
            self._transition("AUDIO")
            try:
                audio_facts = self.audio.transcribe(audio_path)
                audio_text = audio_facts.to_text()
            except Exception as e:
                print(f"audio failed (non-fatal): {e}")

        if video_path:
            self._transition("VIDEO")
            try:
                video_facts = self.video.analyze(video_path)
                video_text = video_facts.to_text()
            except Exception as e:
                print(f"video failed (non-fatal): {e}")

        if document_path:
            self._transition("DOCUMENT")
            try:
                document_text = self.document.extract_text(document_path)
            except Exception as e:
                print(f"document extraction failed (non-fatal): {e}")

        # search decision gate
        self._transition("SEARCH_GATE")
        search_context = ""
        graph_context = ""
        if self._needs_live_data(prompt, force_search, depth):
            self._transition("SEARCH")
            query = prompt[:200]
            try:
                search_context = self.rag.retrieve(query, use_cache=True)
            except Exception as e:
                print(f"search failed (non-fatal): {e}")
            # NER + graph query over live web text + image/audio/video/document text
            self._transition("FACT_EXTRACT")
            try:
                web_facts = self.fact_extractor.triples(search_context) if search_context else []
                for f in web_facts:
                    f.source = "web"
                    self.graph.add_fact(f)
                for text, src in [(image_text, "image"), (audio_text, "audio"),
                                  (video_text, "video"), (document_text, "document")]:
                    if text:
                        for f in self.fact_extractor.triples(text):
                            f.source = src
                            self.graph.add_fact(f)
                graph_context = self.gquery.facts_for_question(self.graph, prompt)
            except Exception as e:
                print(f"fact extraction failed (non-fatal): {e}")

        return {
            "blocked": None,
            "image_facts": image_facts,
            "image_text": image_text,
            "audio_facts": audio_facts,
            "audio_text": audio_text,
            "video_facts": video_facts,
            "video_text": video_text,
            "document_text": document_text,
            "search_context": search_context,
            "graph_context": graph_context,
            "enriched": "\n\n".join(x for x in [image_text, audio_text, video_text, document_text, search_context, graph_context] if x),
        }

    def _build_prompt(self, prompt: str, ctx: dict) -> str:
        if not ctx.get("enriched"):
            return prompt
        return (
            f"[context]\n{ctx['enriched']}\n[/context]\n\n"
            f"Answer using the context above when relevant: {prompt}"
        )

    def _load_model(self):
        config = MorphConfig()
        if self.plugin_dir:
            config.plugin_dir = self.plugin_dir
        self.model = MorphModel(config)

        # base model was loaded bf16 by MorphModel; load the trained adapter
        # + novel components (load_checkpoint wraps base_model_raw in the
        # PEFT adapter itself, so do NOT call apply_lora() here first)
        self.model.load_checkpoint(str(self.model_path))
        self.model.eval()

        if not self.use_cpu and torch.cuda.is_available():
            self.device = torch.device("cuda")
        elif not self.use_cpu and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            self.device = torch.device("mps")
        else:
            self.device = torch.device("cpu")

        self.model = self.model.to(self.device)
        
        # Enable extended context if configured
        if hasattr(self.model.base_model_raw, 'config') and hasattr(self.model.base_model_raw.config, 'rope_scaling'):
            if self.model.base_model_raw.config.rope_scaling:
                print(f"Extended context enabled: {self.model.base_model_raw.config.rope_scaling}")
        
        print(f"Model loaded on {self.device}")
        if self.model._plugins:
            print(f"Loaded plugins: {list(self.model._plugins.keys())}")

    def expand_moe_experts(self):
        """Dynamically expand MoE experts based on usage patterns."""
        if hasattr(self.model, 'moe_block') and hasattr(self.model.moe_block, 'prune_and_expand_experts'):
            self.model.moe_block.prune_and_expand_experts()
            print(f"MoE expanded. Total experts: {len(self.model.moe_block.experts)}")

    def chat(self, prompt: str, max_tokens: int = 512, temperature: float = 0.7,
             skill: Optional[str] = None, image_path: Optional[str] = None,
             audio_path: Optional[str] = None, video_path: Optional[str] = None,
             document_path: Optional[str] = None,
             force_search: bool = False, depth: int = 0,
             use_tools: bool = False) -> str:
        self._transition("INTAKE")
        self._transition("GUARD_IN")

        ctx = self._ingest(prompt, image_path=image_path, audio_path=audio_path,
                          video_path=video_path, document_path=document_path,
                          force_search=force_search, depth=depth)
        if ctx.get("blocked"):
            return ctx["blocked"]

        # L2 routing: explicit skill param wins, else routing matrix
        route = self.routing.route(prompt)
        skill_name = skill if (skill and skill in self.skills) else route.skill
        if skill_name and skill_name in self.skills:
            self._transition("ROUTED")
            self.active_skill = skill_name
            skill_obj = self.skills[skill_name]
            system_prompt = skill_obj.system_prompt if hasattr(skill_obj, "system_prompt") and skill_obj.system_prompt else skill_obj.description
            prompt = self._build_prompt(
                f"system\n{system_prompt}\nuser\n{prompt}\nassistant\n",
                ctx,
            )
        else:
            self.active_skill = None
            prompt = self._build_prompt(f"system\nYou are a helpful assistant.\nuser\n{prompt}\nassistant\n", ctx)
        inputs = self.model.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=self.model.cfg.max_seq_len)
        input_ids = inputs["input_ids"].to(self.device)
        attention_mask = inputs["attention_mask"].to(self.device)

        skill_idx = None
        if self.active_skill:
            skill_idx = route.index if route.skill == self.active_skill else \
                deterministic_index(self.active_skill, self.model.cfg.num_skill_tokens)

        code_feat = build_code_features(self.model.tokenizer, input_ids.cpu()).to(self.device)

        # v6: build multimodal embeddings
        vision_embeds = None
        audio_embeds = None
        video_embeds = None
        if image_path and ctx.get("image_facts") and ctx["image_facts"].embedding is not None:
            vision_embeds = ctx["image_facts"].embedding.to(self.device)
        if audio_path and ctx.get("audio_facts") and ctx["audio_facts"].embedding is not None:
            audio_embeds = ctx["audio_facts"].embedding.to(self.device)
        if video_path and ctx.get("video_facts") and ctx["video_facts"].embeddings is not None:
            video_embeds = ctx["video_facts"].embeddings.to(self.device)

        self._transition("GEN")
        with torch.no_grad():
            outputs = self.model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                skill_token_id=skill_idx,
                code_feat=code_feat,
                max_new_tokens=max_tokens,
                temperature=temperature,
                vision_embeds=vision_embeds,
                audio_embeds=audio_embeds,
                video_embeds=video_embeds,
            )
        self._transition("VERIFY")

        generated_ids = outputs[0][input_ids.shape[1]:]
        response = self.model.tokenizer.decode(generated_ids, skip_special_tokens=True)

        # v6: tool use
        if use_tools:
            self._transition("TOOL_USE")
            tool_response = self._execute_tools(response, prompt)
            if tool_response:
                response = tool_response

        # L5 out-bound guardrails
        self._transition("GUARD_OUT")
        out_dec = self.rules.eval(response, phase="out")
        if out_dec.action == "block":
            response = out_dec.reply or "That output was blocked for safety."
        elif out_dec.action == "mask" and out_dec.masked is not None:
            response = out_dec.masked

        self._persist_turn()
        self.kv.set("last_turn", {"prompt": prompt, "response": response}, ttl=3600)
        self._transition("RESPOND")
        self._transition("IDLE")
        self.active_skill = None
        return response.strip()

    def chat_best_of_n(self, prompt: str, max_tokens: int = 512, n: int = 4,
                       skill: Optional[str] = None, image_path: Optional[str] = None,
                       audio_path: Optional[str] = None, video_path: Optional[str] = None,
                       document_path: Optional[str] = None,
                       force_search: bool = False, depth: int = 0,
                       cross_examine: bool = True, use_tools: bool = False) -> str:
        """Self-critique decoding: generate n candidates, keep the best-scored one."""
        self._transition("INTAKE")
        self._transition("GUARD_IN")

        ctx = self._ingest(prompt, image_path=image_path, audio_path=audio_path,
                          video_path=video_path, document_path=document_path,
                          force_search=force_search, depth=depth)
        if ctx.get("blocked"):
            return ctx["blocked"]

        route = self.routing.route(prompt)
        skill_name = skill if (skill and skill in self.skills) else route.skill
        if skill_name and skill_name in self.skills:
            self._transition("ROUTED")
            self.active_skill = skill_name
            skill_obj = self.skills[skill_name]
            prompt = f"{skill_obj.token}\n{prompt}"
        else:
            self.active_skill = None

        prompt = self._build_prompt(prompt, ctx)
        inputs = self.model.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=self.model.cfg.max_seq_len)
        input_ids = inputs["input_ids"].to(self.device)
        attention_mask = inputs["attention_mask"].to(self.device)

        skill_idx = None
        if self.active_skill:
            skill_idx = route.index if route.skill == self.active_skill else \
                deterministic_index(self.active_skill, self.model.cfg.num_skill_tokens)

        code_feat = build_code_features(self.model.tokenizer, input_ids.cpu()).to(self.device)

        # v6: multimodal embeddings
        vision_embeds = None
        audio_embeds = None
        video_embeds = None
        if image_path and ctx.get("image_facts") and ctx["image_facts"].embedding is not None:
            vision_embeds = ctx["image_facts"].embedding.to(self.device)
        if audio_path and ctx.get("audio_facts") and ctx["audio_facts"].embedding is not None:
            audio_embeds = ctx["audio_facts"].embedding.to(self.device)
        if video_path and ctx.get("video_facts") and ctx["video_facts"].embeddings is not None:
            video_embeds = ctx["video_facts"].embeddings.to(self.device)

        self._transition("GEN")
        with torch.no_grad():
            outputs = self.model.generate_best_of_n(
                input_ids=input_ids,
                attention_mask=attention_mask,
                skill_token_id=skill_idx,
                code_feat=code_feat,
                n=n,
                max_new_tokens=max_tokens,
                accept_threshold=0.6,
                early_exit_margin=0.01,
                vision_embeds=vision_embeds,
                audio_embeds=audio_embeds,
                video_embeds=video_embeds,
            )
        self._transition("VERIFY")

        generated_ids = outputs[0][input_ids.shape[1]:]
        response = self.model.tokenizer.decode(generated_ids, skip_special_tokens=True)

        # cross-examination
        if cross_examine:
            response = self._cross_examine(response, ctx)

        # v6: tool use
        if use_tools:
            self._transition("TOOL_USE")
            tool_response = self._execute_tools(response, prompt)
            if tool_response:
                response = tool_response

        self._transition("GUARD_OUT")
        out_dec = self.rules.eval(response, phase="out")
        if out_dec.action == "block":
            response = out_dec.reply or "That output was blocked for safety."
        elif out_dec.action == "mask" and out_dec.masked is not None:
            response = out_dec.masked

        self._persist_turn()
        self.kv.set("last_turn", {"prompt": prompt, "response": response}, ttl=3600)
        self._transition("RESPOND")
        self._transition("IDLE")
        self.active_skill = None
        return response.strip()

    def _cross_examine(self, response: str, ctx: dict) -> str:
        """Score the best-of-n winner against image + web facts. If entity
        overlap with grounded facts is low and the prompt was fact-dependent,
        append a self-correction note (does not silently fabricate)."""
        ground_truth = (ctx.get("image_text") or "") + " " + (ctx.get("graph_context") or "")
        if not ground_truth.strip():
            return response
        resp_ents = set()
        for kind in ("PERSON", "ORG", "LOCATION", "DATE", "NUMBER"):
            resp_ents |= self.fact_extractor.extract(response).get(kind, set())
        truth_ents = set()
        for kind in ("PERSON", "ORG", "LOCATION", "DATE", "NUMBER"):
            truth_ents |= self.fact_extractor.extract(ground_truth).get(kind, set())
        if not truth_ents:
            return response
        overlap = len(resp_ents & truth_ents) / len(truth_ents)
        out = self.rules.eval(response, phase="out")
        compliant = out.action != "block"
        print(f"cross-examine: entity overlap={overlap:.2f}, compliant={compliant}")
        if overlap < 0.5 and compliant:
            return response + (
                "\n\n[verifier] This answer only weakly overlaps the retrieved "
                "facts; treat details against the cited context above."
            )
        return response

    def _execute_tools(self, response: str, original_prompt: str) -> Optional[str]:
        """Extract and execute tool calls from model response if present."""
        try:
            tool_calls = self.tool_registry.parse_calls(response)
            if not tool_calls:
                return None
            results = []
            for call in tool_calls:
                result = self.tool_registry.execute(call)
                results.append(f"[tool:{call.tool_name}] {result}")
            return "\n".join(results)
        except Exception as e:
            print(f"tool execution failed: {e}")
            return None

    def _persist_turn(self):
        """Write the last refined hidden state into the cross-turn scratchpad
        and persistent memory so later turns can condition on it."""
        refined = getattr(self.model, "_last_refined", None)
        if refined is None:
            return
        try:
            if refined.dim() == 3:
                self.model.scratchpad.write(refined)
                self.model.memory.write(refined)
        except Exception as e:
            print(f"persist_turn failed (non-fatal): {e}")

    def install_skill(self, skill_path: str) -> bool:
        skill = Skill.load(skill_path)
        if skill.name in self.skills:
            print(f"Skill '{skill.name}' already installed")
            return False

        self.skills[skill.name] = skill
        self.routing.load_legacy_skill(skill_path)
        print(f"Skill '{skill.name}' installed: {skill.description}")
        print(f"  Trigger patterns: {', '.join(skill.trigger_patterns)}")
        return True

    def uninstall_skill(self, skill_name: str) -> bool:
        if skill_name in self.skills:
            del self.skills[skill_name]
            if self.active_skill == skill_name:
                self.active_skill = None
            print(f"Skill '{skill_name}' uninstalled")
            return True
        return False

    def list_skills(self) -> List[Dict[str, Union[str, List[str], bool]]]:
        return [
            {
                "name": s.name,
                "description": s.description,
                "token": s.token,
                "patterns": s.trigger_patterns,
                "installed": s.installed
            }
            for s in self.skills.values()
        ]

    def auto_route(self, prompt: str) -> Optional[str]:
        return self.routing.route(prompt).skill

    def chat_auto(self, prompt: str, max_tokens: int = 512, temperature: float = 0.7) -> str:
        skill = self.auto_route(prompt)
        return self.chat(prompt, max_tokens, temperature, skill)

    def create_skill_gguf(self, skill_name: str, output_path: str):
        if not HAS_LLAMACPP:
            raise ImportError("llama-cpp-python not installed. Run: pip install llama-cpp-python")

        skill = self.skills.get(skill_name)
        if not skill or not skill.lora_weights:
            raise ValueError(f"Skill '{skill_name}' not found or has no weights")

        lora_path = f"{output_path}/{skill_name}_lora.bin"
        with open(lora_path, 'wb') as f:
            f.write(skill.lora_weights)

        gguf_model = Llama(
            model_path=str(self.model_path / "model.gguf"),
            lora_path=lora_path,
            n_ctx=self.model.cfg.max_seq_len,
            n_gpu_layers=-1 if torch.cuda.is_available() else 0
        )
        print(f"GGUF model with skill '{skill_name}' created at {output_path}")

    def export_to_gguf(self, output_dir: str):
        if not HAS_LLAMACPP:
            raise ImportError("llama-cpp-python not installed")

        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)

        print("Converting to GGUF format...")
        import subprocess
        result = subprocess.run([
            "python", "-m", "transformers.convert_save_to_hf",
            "--input_path", str(self.model_path / "base_lora"),
            "--output_path", str(output_path / "hf_model")
        ], capture_output=True, text=True)

        print(f"GGUF export initiated. Use llama.cpp convert script for final conversion.")
        print(f"Output directory: {output_dir}")


def load_model(model_path: str, **kwargs) -> MorphRuntime:
    return MorphRuntime(model_path, **kwargs)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="MORPH-AI local runtime")
    parser.add_argument("--model", required=True, help="Path to trained model folder")
    parser.add_argument("--no-4bit", action="store_true", help="Disable 4-bit quantization")
    parser.add_argument("--cpu", action="store_true", help="Force CPU inference")
    args = parser.parse_args()

    rt = MorphRuntime(args.model, use_4bit=not args.no_4bit, use_cpu=args.cpu)
    print("\nMORPH-AI ready. Type 'exit' to quit.")
    print(f"Installed skills: {[s.name for s in rt.skills.values()]}")
    while True:
        try:
            prompt = input("\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if not prompt:
            continue
        if prompt.lower() in ("exit", "quit"):
            break
        print(rt.chat_auto(prompt, max_tokens=512))