"""
AEC AI Reader - Inference Engine
Mid-2026 Two-Stage Generation Architecture
"""
import json
import re
from typing import Dict, Any, Generator, Optional
from llama_cpp import Llama
from llama_cpp.llama_grammar import LlamaGrammar
from .chain_cache import ChainCache

class AECInferenceEngine:
    def __init__(
        self,
        model_path: str,
        grammar_path: str = "serving/grammar.gbnf",
        cache_db_path: str = "chain_cache.sqlite",
        n_ctx: int = 2048,
        n_threads: Optional[int] = None
    ):
        print(f"[Inference] Loading model: {model_path}")
        self.llm = Llama(
            model_path=model_path,
            n_ctx=n_ctx,
            n_threads=n_threads,
            verbose=False
        )
        
        print(f"[Inference] Loading GBNF grammar: {grammar_path}")
        with open(grammar_path, "r") as f:
            grammar_text = f.read()
        self.grammar = LlamaGrammar.from_string(grammar_text)
        
        print("[Inference] Connecting to Chain Cache...")
        self.cache = ChainCache(db_path=cache_db_path)
        
        self.system_prompt = (
            "Kamu adalah arsitek AI AEC. "
            "Baca instruksi dan hasilkan JSON untuk pascal-editor MCP.\n"
            "Gunakan tag <think>...</think> untuk penalaran logis dan spasial, "
            "lalu berikan output akhir murni dalam format JSON."
        )

    def process_instruction(self, instruction: str) -> Dict[str, Any]:
        """
        End-to-end processing: Cache lookup -> Two-Stage Inference (if miss).
        """
        # 1. Cache Lookup (Semantic/Exact)
        cache_result = self.cache.lookup(instruction)
        if cache_result:
            return {
                "source": "cache",
                "similarity": cache_result["similarity"],
                "thinking": cache_result.get("thinking", ""),
                "output_type": cache_result["output_type"],
                "output": cache_result["output"]
            }

        # 2. Cache Miss -> LLM Inference (Two-Stage)
        prompt = f"<|im_start|>system\n{self.system_prompt}<|im_end|>\n<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n<think>\n"
        
        # Stage 1: Free generation for thinking process
        thinking_text = ""
        stage1_res = self.llm(
            prompt,
            max_tokens=512,
            stop=["</think>"],
            stream=False
        )
        
        thinking_text = stage1_res["choices"][0]["text"].strip()
        
        # Append the thinking part to prompt for stage 2
        prompt_stage2 = prompt + thinking_text + "\n</think>\n"
        
        # Stage 2: Constrained generation for JSON output
        stage2_res = self.llm(
            prompt_stage2,
            max_tokens=1024,
            grammar=self.grammar,
            stop=["<|im_end|>"],
            stream=False
        )
        
        json_output_str = stage2_res["choices"][0]["text"].strip()
        
        try:
            parsed_json = json.loads(json_output_str)
        except json.JSONDecodeError:
            parsed_json = {"error": "Invalid JSON generated", "raw": json_output_str}
            
        output_type = parsed_json.get("output_type", "unknown")
        
        # 3. Cache the new result
        if "error" not in parsed_json:
            self.cache.add(
                instruction=instruction,
                output=parsed_json.get("output", {}),
                output_type=output_type,
                thinking=thinking_text
            )

        return {
            "source": "llm",
            "thinking": thinking_text,
            "output_type": output_type,
            "output": parsed_json.get("output", {})
        }

if __name__ == "__main__":
    # Test script setup
    import sys
    model_path = sys.argv[1] if len(sys.argv) > 1 else "../models/qwen3-4b-instruct-q4_k_m.gguf"
    print(f"To test, run: python -m serving.inference <model_path>")