#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
==============================================================================
Manara-3B | Inference & Evaluation Script
==============================================================================
This script loads the fine-tuned Manara-3B model and runs inference,
including the self-reflective critic pass for guardrail compliance.

Usage:
  python inference.py \
      --model_path ../../model_weights/manara-3b-final \
      --prompt "ما هي متطلبات نسبة كفاية رأس المال وفق تعليمات بنك الكويت المركزي؟" \
      --language ar
==============================================================================
"""

import sys
import argparse
import logging
from pathlib import Path

import torch

try:
    from unsloth import FastLanguageModel
    UNSLOTH_AVAILABLE = True
except ImportError:
    UNSLOTH_AVAILABLE = False
    from transformers import AutoModelForCausalLM, AutoTokenizer

# Add guardrails to path
sys.path.append(str(Path(__file__).parent.parent / "guardrails"))
from self_reflective_loop import ManaraGuardrailEngine

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

ALPACA_PROMPT_AR = """فيما يلي تعليمات تصف مهمة. اكتب استجابة تكمل الطلب بشكل مناسب.

### التعليمات:
{}

### الاستجابة:
"""

ALPACA_PROMPT_EN = """Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{}

### Response:
"""


class ManaraInferenceEngine:
    """Handles inference for the Manara-3B model with guardrail integration."""

    def __init__(self, model_path: str, max_seq_length: int = 2048):
        self.model_path = model_path
        self.max_seq_length = max_seq_length
        self._load_model()
        self.guardrail_engine = ManaraGuardrailEngine()

    def _load_model(self):
        """Loads the fine-tuned model for inference."""
        if UNSLOTH_AVAILABLE:
            self.model, self.tokenizer = FastLanguageModel.from_pretrained(
                model_name=self.model_path,
                max_seq_length=self.max_seq_length,
                dtype=None,
                load_in_4bit=True,
            )
            FastLanguageModel.for_inference(self.model)  # Enable native 2x faster inference
        else:
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_path,
                device_map="auto",
                torch_dtype=torch.float16,
            )
        logger.info(f"Model loaded from: {self.model_path}")

    def generate(
        self,
        prompt: str,
        language: str = "ar",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
        top_p: float = 0.9,
        apply_guardrails: bool = True,
    ) -> dict:
        """
        Generates a response for the given prompt, with optional guardrail checking.

        Returns a dict containing:
          - 'response': The generated text
          - 'guardrail_passed': Whether the response passed all guardrails
          - 'corrections': Any corrections made by the critic pass
        """
        # Format prompt
        template = ALPACA_PROMPT_AR if language == "ar" else ALPACA_PROMPT_EN
        formatted_prompt = template.format(prompt)

        # Tokenize
        inputs = self.tokenizer([formatted_prompt], return_tensors="pt").to("cuda")

        # Generate
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
            )

        # Decode (only the new tokens)
        generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
        response = self.tokenizer.decode(generated_ids, skip_special_tokens=True)

        result = {
            "prompt": prompt,
            "response": response,
            "language": language,
            "guardrail_passed": True,
            "corrections": [],
        }

        # Apply self-reflective guardrail check
        if apply_guardrails:
            guardrail_result = self.guardrail_engine.check(response, language)
            result["guardrail_passed"] = guardrail_result["passed"]
            result["corrections"] = guardrail_result["violations"]

            if not guardrail_result["passed"]:
                logger.warning(f"Guardrail violations detected: {guardrail_result['violations']}")
                # Optionally regenerate with corrective prompt
                # result["response"] = self._regenerate_with_correction(prompt, response, guardrail_result)

        return result


def parse_args():
    parser = argparse.ArgumentParser(description="Manara-3B: Inference")
    parser.add_argument("--model_path", required=True, help="Path to fine-tuned model")
    parser.add_argument("--prompt", required=True, help="Input prompt")
    parser.add_argument("--language", default="ar", choices=["ar", "en"])
    parser.add_argument("--max_new_tokens", type=int, default=512)
    parser.add_argument("--no_guardrails", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    engine = ManaraInferenceEngine(args.model_path)
    result = engine.generate(
        prompt=args.prompt,
        language=args.language,
        max_new_tokens=args.max_new_tokens,
        apply_guardrails=not args.no_guardrails,
    )
    print("\n" + "=" * 60)
    print(f"PROMPT: {result['prompt']}")
    print("=" * 60)
    print(f"RESPONSE:\n{result['response']}")
    print("=" * 60)
    print(f"GUARDRAIL PASSED: {result['guardrail_passed']}")
    if result["corrections"]:
        print(f"VIOLATIONS: {result['corrections']}")
    print("=" * 60)