"""
Evaluation script for the IaC-Coder model.
Tests the model on representative IaC tasks across Terraform, Ansible, Docker, K8s, and Shell.

Usage:
  pip install transformers torch peft
  python eval_iac.py
"""

import torch
from transformers import pipeline

MODEL_ID = "Tejas86/iac-coder-1.5b"

print(f"Loading model from {MODEL_ID}...")
pipe = pipeline(
    "text-generation",
    model=MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

test_cases = [
    {
        "name": "Terraform: Create an S3 bucket with versioning",
        "messages": [
            {"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready Terraform HCL code."},
            {"role": "user", "content": "Task: Create an AWS S3 bucket with versioning enabled and server-side encryption\n\nGenerate the Terraform HCL code."},
        ],
    },
    {
        "name": "Terraform: Modify VPC to add private subnet",
        "messages": [
            {"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready Terraform HCL code."},
            {"role": "user", "content": 'Task: Add a private subnet to the VPC\n\nModify the following Terraform HCL file:\n```\nresource "aws_vpc" "main" {\n  cidr_block = "10.0.0.0/16"\n  tags = {\n    Name = "main-vpc"\n  }\n}\n\nresource "aws_subnet" "public" {\n  vpc_id     = aws_vpc.main.id\n  cidr_block = "10.0.1.0/24"\n  tags = {\n    Name = "public-subnet"\n  }\n}\n```'},
        ],
    },
    {
        "name": "Ansible: Install and configure Nginx",
        "messages": [
            {"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready YAML (Ansible/K8s) code."},
            {"role": "user", "content": "Task: Write an Ansible playbook to install nginx on Ubuntu and start the service\n\nGenerate the YAML (Ansible/K8s) code."},
        ],
    },
    {
        "name": "Kubernetes: Redis StatefulSet with PVC",
        "messages": [
            {"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready YAML (Ansible/K8s) code."},
            {"role": "user", "content": "Task: Create a Kubernetes StatefulSet for Redis with persistent storage\n\nGenerate the YAML (Ansible/K8s) code."},
        ],
    },
    {
        "name": "Dockerfile: Multi-stage Python Flask app",
        "messages": [
            {"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready Dockerfile code."},
            {"role": "user", "content": "Task: Create a multi-stage Dockerfile for a Python Flask application with pip dependencies\n\nGenerate the Dockerfile code."},
        ],
    },
    {
        "name": "Shell: CI/CD deployment script",
        "messages": [
            {"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready Shell/Bash script code."},
            {"role": "user", "content": "Task: Write a bash script that builds a Docker image, runs tests, and pushes to ECR\n\nGenerate the Shell/Bash script code."},
        ],
    },
]

print("\n" + "=" * 80)
print("IaC-CODER MODEL EVALUATION")
print("=" * 80)

for i, test in enumerate(test_cases, 1):
    print(f"\n{'─' * 80}")
    print(f"Test {i}/{len(test_cases)}: {test['name']}")
    print(f"{'─' * 80}")
    
    output = pipe(
        test["messages"],
        max_new_tokens=512,
        temperature=0.2,
        do_sample=True,
        top_p=0.9,
    )
    
    generated = output[0]["generated_text"]
    if isinstance(generated, list):
        assistant_msgs = [m for m in generated if m["role"] == "assistant"]
        response = assistant_msgs[-1]["content"] if assistant_msgs else str(generated)
    else:
        response = generated
    
    print(f"\nGenerated:\n{response[:1500]}")

print(f"\n{'=' * 80}")
print(f"Evaluation complete! Model: {MODEL_ID}")
print(f"{'=' * 80}")