import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import json
import os

# For internal use: base model details are in PRIVATE_MODEL_TRAINING_NOTES.md
def load_maeyen_assistant(model_path):
    """Load Maeyen Trust & Risk Assistant (internal use only)"""
    # See PRIVATE_MODEL_TRAINING_NOTES.md for base model info
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    tokenizer.pad_token = tokenizer.eos_token
    
    print("Loading model...")
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        device_map="cpu"  # or "auto" if GPU available
    )
    
    return model, tokenizer

def generate_recommendation(model, tokenizer, task, data):
    """Generate structured recommendation from Maeyen AI"""
    system_prompts = {
        "risk": "You are Maeyen AI Transaction Risk Agent. Assess risk and output valid JSON only with requires_human_review: true.",
        "evidence": "You are Maeyen AI Evidence Review Agent. Review evidence and output valid JSON only with requires_human_review: true.",
        "dispute": "You are Maeyen AI Dispute Assistant. Summarize dispute and output valid JSON only with requires_human_review: true.",
        "trust": "You are Maeyen AI Trust Score Explanation Agent. Explain trust score and output valid JSON only."
    }

    system_prompt = system_prompts.get(task, system_prompts["risk"])
    prompt = f"""<|im_start|>system
{system_prompt}<|im_end|>
<|im_start|>user
{json.dumps(data, indent=2)}<|im_end|>
<|im_start|>assistant
"""

    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.1,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    try:
        return json.loads(response)
    except json.JSONDecodeError:
        return {"raw_response": response, "requires_human_review": True}

# Example usage (internal only)
if __name__ == "__main__":
    # Note: See PRIVATE_MODEL_TRAINING_NOTES.md for full setup
    print("This is for internal use. See PRIVATE_MODEL_TRAINING_NOTES.md for base model details.")