axis/axis_core/controller.py

from __future__ import annotations

import gc
import uuid
import json
import re
import os
import time
from typing import Any, Dict, List, Optional, Generator

from .db import LatticeDB
from .miner import OnDemandMiner
from .decompose import DecompositionLayer
from .match import PatternMatchingLayer
from .simulate import SimulationLayer
from .renderer import GrammarLayer, RenderContext
from .memory import MemoryLattice
from .history import HistoryManager
from .util import now_hms
from .events import EventRecorder

class AXISController:
    """
    AXIS V13.15 Controller: Sovereign Session Edition
    - Manages reasoning stages (0/5 to 5/5) within a specific Cross session.
    """

    def __init__(self, db: LatticeDB, miner: Optional[OnDemandMiner] = None):
        self.db = db
        self.miner = miner
        self.layer_decompose = DecompositionLayer()
        self.layer_match = PatternMatchingLayer()
        self.layer_simulate = SimulationLayer()
        self.layer_grammar = GrammarLayer()
        storage_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data", "memory")
        self.mll = MemoryLattice(os.path.join(storage_dir, "global_memory.json"))
        self.history = HistoryManager(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data", "history"))

    def _sanitize_mined_rules(self, rules: List[Dict[str, Any]], query: str) -> List[Dict[str, Any]]:
        """
        Quarantine mined rules to prevent noise injection.
        """
        valid_domains = {"modal_logic", "topology", "logic", "set_theory", "general"}
        valid_prefixes = ("frame:", "condition:", "logic:")
        sanitized = []
        
        for r in rules:
            if r.get("domain", "").lower() not in valid_domains: continue
            if not all(k in r for k in ["id", "type", "requires", "evidence"]): continue
            
            reqs = r.get("requires", [])
            if not reqs or not isinstance(reqs, list): continue
            if not all(any(req.startswith(p) for p in valid_prefixes) for req in reqs): continue
            
            entails = r.get("entails", [])
            match_all = r.get("match_all", [])
            targets = entails + match_all
            if not targets: continue
            
            valid_targets = [t for t in targets if len(t) >= 3 and not t.isdigit()]
            if not valid_targets: continue
            
            ev = r.get("evidence", 0)
            if r.get("type") == "counter_example":
                ev = -10 
            else:
                ev = max(1, min(3, int(ev)))
            r["evidence"] = ev
            sanitized.append(r)
            
        return sanitized

    def _dump_mined_rules(self, rules: List[Dict[str, Any]], query: str):
        """
        Append mined rules to a dump file for manual review and promotion.
        """
        dump_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data", "mined_rules_dump.jsonl")
        try:
            entry = {
                "ts": time.time(),
                "query": query,
                "rules": rules
            }
            with open(dump_path, "a", encoding="utf-8") as f:
                f.write(json.dumps(entry, ensure_ascii=False) + "\n")
        except Exception as e:
            print(f"⚠️ Failed to dump mined rules: {e}")

    def handle_stream(self, query: str, session_id: Optional[str] = None) -> Generator[Dict[str, Any], None, None]:
        start_time = time.time()
        sid = session_id or self.history.create_session()
        cross_data = self.history.load_session(sid)
        
        if cross_data is None:
            cross_data = {
                "session_id": sid,
                "local_defs": {},
                "messages": [],
                "state": {"stage": "IDLE", "progress": "0/5"}
            }
        
        def emit(stage: str, event_type: str, payload: dict = None, progress: str = "0/5"):
            return {
                "t": round(time.time() - start_time, 3),
                "stage": stage,
                "type": event_type,
                "progress": progress,
                "payload": payload or {}
            }

        self._full_reset()
        yield emit("SESSION_INIT", "STAGE_START", progress="0/5")

        try:
            # --- L1: DECOMPOSE (1/5) ---
            yield emit("DECOMPOSE", "STAGE_START", progress="1/5")
            decomp = self.layer_decompose.decompose(query)
            decomps = [decomp]
            yield emit("DECOMPOSE", "SLOT_EXTRACTED", {"choices_count": len(decomp.slots.get("choices", {}))}, progress="1/5")
            
            # --- V13.10: Interactive Confirmation Check ---
            # If ANY block needs confirmation, we pause the whole stream
            for d in decomps:
                print(f"DEBUG: Checking confirmation for block. Meta: {d.meta}")
                if d.meta.get("needs_confirmation"):
                    repair_data = d.meta["repair_data"]
                    yield emit("PASTE_REPAIR", "USER_CONFIRM_REQUIRED", {
                        "msg": "入力形式が崩れている可能性があります。以下の整形案で解析しますか？",
                        "pretty": repair_data["pretty"],
                        "options": ["yes", "no"]
                    }, progress="1/5")
                    
                    # Persist state for next turn
                    cross_data["state"] = {
                        "stage": "PENDING_CONFIRMATION", 
                        "repair_data": repair_data,
                        "original_query": query
                    }
                    self.history.save_session(sid, cross_data)
                    return # STOP processing here

            combined_responses = []
            
            # --- L2: MINING / RULE LOOKUP (2/5) ---
            yield emit("MINE", "STAGE_START", progress="2/5")
            
            mined_rules = []
            mined_defs = cross_data.get("local_defs", {})
            all_ops = []
            is_provisional = False
            
            if self.miner:
                k_res = self.miner.mine_knowledge(decomp.entities)
                if "DEFINITIONS" in k_res:
                    for d in k_res["DEFINITIONS"]:
                        mined_defs[d["term"]] = d["definition"]
            
            block_text = decomp.text_part or query
            block_reqs = self._derive_requirements(decomp)
            for block_name in block_reqs:
                ops = self._get_fallback_ops(block_name)
                all_ops.extend(ops)

            # --- L3: SIMULATE (3/5) ---
            yield emit("SIMULATE", "STAGE_START", progress="3/5")
            yield emit("SIMULATE", "ELECTRIC_PULSE_START")
            
            # --- PHASE 1: Try with Rules DB Only ---
            yield emit("SIMULATE", "LOG_UPDATE", {"msg": "Phase 1: DB Verification"}, progress="3/5")
            knowledge_p1 = {"OPS": all_ops, "DEFINITIONS": mined_defs, "MINED_RULES": [], "QUERY_TEXT": block_text}
            sim_content = self.layer_simulate.simulate(None, decomp, knowledge_p1)
            
            sim_meta = sim_content.get("meta", {})
            evidence_count = sim_meta.get("evidence_count", 0)
            final_verdict = sim_meta.get("final_verdict", "accepted")

            # --- PHASE 2: Mining Fallback ---
            if final_verdict == "insufficient_evidence" and self.miner:
                yield emit("SIMULATE", "LOG_UPDATE", {"msg": "Phase 2: Mining Fallback Initiated"}, progress="3/5")
                yield emit("MINE", "MINING_FLOW_START", {"target": "FALLBACK_RULES"}, progress="2/5")
                
                r_res = self.miner.mine_formal_rules(entities=decomp.entities, domain="Logic")
                if "STRUCTURE" in r_res:
                    raw_rules = r_res["STRUCTURE"].get("rules", [])
                    mined_rules = self._sanitize_mined_rules(raw_rules, block_text)
                
                if mined_rules:
                    is_provisional = True
                    self._dump_mined_rules(mined_rules, block_text)
                    yield emit("MINE", "RULES_MINED", {"count": len(mined_rules), "sanitized": True}, progress="2/5")
                    knowledge_p2 = {"OPS": all_ops, "DEFINITIONS": mined_defs, "MINED_RULES": mined_rules, "QUERY_TEXT": block_text}
                    sim_content = self.layer_simulate.simulate(None, decomp, knowledge_p2)
                    self.miner.reject()
                    yield emit("MINE", "MEMORY_REJECTED", progress="2/5")
                else:
                    yield emit("MINE", "NO_RULES_MINED", progress="2/5")
            else:
                yield emit("MINE", "SKIPPED_DB_HIT", {"evidence": evidence_count}, progress="2/5")
            
            for step in sim_content.get("logic_check", []):
                yield emit("SIMULATE", "LOG_UPDATE", {"msg": step}, progress="3/5")

            yield emit("SIMULATE", "ELECTRIC_PULSE_END")
            
            # --- L4: RENDER (4/5) ---
            yield emit("RENDER", "STAGE_START", progress="4/5")
            render_ctx = RenderContext(intent=decomp.intent, content=sim_content)
            render_res = self.layer_grammar.render(render_ctx)
            
            # --- L5: COMMIT & PERSIST (5/5) ---
            stage_name = "COMMIT"
            yield emit(stage_name, "STAGE_START", progress="5/5")
            
            final_text = render_res["text"]
            if is_provisional:
                final_text += "\n\n(Note: This result includes provisional findings from dynamic rule mining.)"
            
            all_entities = decomp.entities
            verified_facts = [l for l in sim_content.get("answer", "").split("\n") if l.strip()]
            self.mll.add_turn(sid, query, final_text, all_entities, verified_facts)
            self.mll.save()

            cross_data["messages"].append({"role": "user", "content": query, "ts": time.time()})
            cross_data["messages"].append({"role": "assistant", "content": final_text, "ts": time.time()})
            cross_data["local_defs"] = mined_defs 
            cross_data["state"] = {"stage": "DONE", "progress": "5/5"}
            cross_data["query"] = query 
            self.history.save_session(sid, cross_data)

            yield emit("SUCCESS", "RESULT_FINAL", {
                "response": final_text,
                "attainment": sim_content.get("attainment"),
                "meta": {"provisional": is_provisional}
            }, progress="5/5")

        except Exception as e:
            import traceback
            traceback.print_exc()
            yield emit("FAILURE", "ERROR_THROWN", {"msg": str(e)})
        finally:
            self._full_reset()

    def _derive_requirements(self, decomp: DecompositionResult) -> List[str]:
        if decomp.slots.get("choices"):
            return ["BLOCK_CHOICE_SOLVER"]
        
        # Expanded intent check for Explanation Generator
        explanation_intents = {
            "explain", "describe_features", "how_to", "compare", "why", 
            "pros_cons", "examples", "summary", "history", "define",
            "assumption_analysis", "scope_and_limits", "failure_modes", "why_not"
        }
        if decomp.intent in explanation_intents:
            return ["BLOCK_EXPLANATION_GENERATOR"]
            
        return ["BLOCK_GENERAL_INFERENCE"]

    def _get_fallback_ops(self, block_name: str) -> List[Dict[str, Any]]:
        fallback_map = {
            "BLOCK_CHOICE_SOLVER": [
                {"op": "EXTRACT_CLAIMS", "args": {}},
                {"op": "REGISTER_DEFS", "args": {"source": "mined"}},
                {"op": "LEXICAL_MATCH", "args": {}},
                {"op": "VERIFY_WITH_RULESET", "args": {}},
                {"op": "VERIFY_NECESSITY", "args": {}},
                {"op": "SOLVE_CHOICE", "args": {}}
            ],
            "BLOCK_EXPLANATION_GENERATOR": [
                {"op": "REGISTER_DEFS", "args": {"source": "mined"}},
                # {"op": "APPLY_SCHEMA", "args": {}}, # If schema exists
                {"op": "GENERATE_EXPLANATION", "args": {}}
            ],
            "BLOCK_GENERAL_INFERENCE": [
                {"op": "REGISTER_DEFS", "args": {"source": "mined"}},
                {"op": "LEXICAL_MATCH", "args": {}},
            ]
        }
        return fallback_map.get(block_name, [{"op": "LEXICAL_MATCH", "args": {}}])

    def _full_reset(self) -> None:
        if self.miner: self.miner.reject()
        gc.collect()
        try:
            import torch
            if torch.cuda.is_available(): torch.cuda.empty_cache()
            elif torch.backends.mps.is_available(): torch.mps.empty_cache()
        except: pass