#!/usr/bin/env python3
"""
Inference script for trained LoRA adapter.

Loads the trained model and generates Cyron summaries.
"""

import argparse
from pathlib import Path


def load_model(model_path):
    """Load trained LoRA model."""
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer
    
    print(f"Loading base model from {model_path}...")
    base_model = AutoModelForCausalLM.from_pretrained(
        model_path,
        device_map="auto",
        torch_dtype="auto",
    )
    
    print(f"Loading LoRA weights from {model_path}...")
    model = PeftModel.from_pretrained(base_model, model_path)
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    
    return model, tokenizer


def generate_summary(model, tokenizer, task, files_changed=None, tests_run=False, test_count=0, commit=False, push=False, errors_seen=False):
    """Generate a Cyron summary for a task."""
    
    # Build prompt
    prompt = f"Task: {task}"
    if files_changed:
        prompt += f"\nFiles: {', '.join(files_changed[:3])}"
    
    if tests_run:
        prompt += f"\nTests run: {test_count}"
    if commit:
        prompt += "\nCommit: yes"
    if push:
        prompt += "\nPush: yes"
    if errors_seen:
        prompt += "\nErrors seen: yes"
    
    prompt += "\n\nGenerate a Cyron summary:"
    
    # Tokenize
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # Generate
    outputs = model.generate(
        **inputs,
        max_new_tokens=500,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
    )
    
    # Decode
    generated = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    
    return generated


def main():
    parser = argparse.ArgumentParser(description="Generate Cyron summaries with LoRA")
    parser.add_argument("--model", type=str, default="output/ornith-35b-lora-cyron", 
                       help="Path to trained LoRA model")
    parser.add_argument("--task", type=str, required=True, help="Task description")
    parser.add_argument("--files", type=str, nargs="*", default=[], help="Files changed")
    parser.add_argument("--tests-run", action="store_true", help="Tests were run")
    parser.add_argument("--test-count", type=int, default=0, help="Number of tests")
    parser.add_argument("--commit", action="store_true", help="Commit was made")
    parser.add_argument("--push", action="store_true", help="Code was pushed")
    parser.add_argument("--errors", action="store_true", help="Errors were seen")
    args = parser.parse_args()
    
    model, tokenizer = load_model(args.model)
    
    summary = generate_summary(
        model,
        tokenizer,
        task=args.task,
        files_changed=args.files,
        tests_run=args.tests_run,
        test_count=args.test_count,
        commit=args.commit,
        push=args.push,
        errors_seen=args.errors,
    )
    
    print("\nGenerated Summary:")
    print("=" * 70)
    print(summary)
    print("=" * 70)


if __name__ == "__main__":
    main()