#!/usr/bin/env python3
"""
BRITISH PROMPT OPTIMIZER - PRODUCTION VERSION
Aceita QUALQUER idioma e SEMPRE retorna prompt otimizado em inglês britânico
"""

from ollama import Client
import json
import logging
from datetime import datetime
from pathlib import Path

# Configuração
OLLAMA_API_KEY = "08d341f2a47745999691ebbde61a3374.cz1ftaVA-TViLz3vwGZBAFLm"
MODEL = "gpt-oss:20b"

# Sistema para OTIMIZAÇÃO DE PROMPTS
SYSTEM_PROMPT = """You are a PROMPT OPTIMIZER for a British International School.

YOUR TASK:
1. Take ANY input (Portuguese, American English, any language)
2. ALWAYS output an OPTIMIZED PROMPT in British English
3. The output should be a clear, educational prompt suitable for British school levels

OPTIMIZATION RULES:
- Convert to British spelling: colour, centre, analyse, organise, behaviour, favourite
- Add British educational context: Pre-Prep, Prep, Junior, IGCSE, IBDP
- Make the prompt clear and specific
- Include learning objectives when relevant
- Use British pedagogical terminology

EXAMPLES:
Input: "como ensinar matemática?"
Output: "Design a mathematics lesson plan for Year 5 pupils covering fractions and decimals, incorporating visual aids and hands-on activities aligned with the National Curriculum."

Input: "teach photosynthesis"
Output: "Create an engaging IGCSE Biology lesson on photosynthesis, including practical experiments, diagrams, and assessment strategies suitable for Year 10 pupils."

ALWAYS OUTPUT ONLY THE OPTIMIZED PROMPT, NOTHING ELSE."""

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class BritishPromptOptimizer:
    def __init__(self):
        self.client = Client(
            host="https://ollama.com",
            headers={'Authorization': OLLAMA_API_KEY}
        )
        self.model = MODEL
        
    def optimize_prompt(self, user_input):
        """Optimize ANY input to British English educational prompt"""
        
        try:
            # Generate optimized prompt
            response = self.client.chat(
                model=self.model,
                messages=[
                    {'role': 'system', 'content': SYSTEM_PROMPT},
                    {'role': 'user', 'content': user_input}
                ],
                options={
                    'temperature': 0.7,
                    'top_p': 0.9,
                    'max_tokens': 200
                }
            )
            
            optimized = response['message']['content'].strip()
            
            # Ensure British spelling
            optimized = self.enforce_british_spelling(optimized)
            
            return optimized
            
        except Exception as e:
            logger.error(f"Error optimizing prompt: {e}")
            return "Create an engaging lesson plan for British pupils incorporating active learning strategies."
    
    def enforce_british_spelling(self, text):
        """Ensure British spelling in output"""
        replacements = {
            "color": "colour",
            "center": "centre",
            "analyze": "analyse",
            "organize": "organise",
            "behavior": "behaviour",
            "favor": "favour",
            "honor": "honour",
            "theater": "theatre",
            "meter": "metre",
            "fiber": "fibre",
            "defense": "defence",
            "license": "licence",
            "practice" + " (noun)": "practice",
            "practise" + " (verb)": "practise"
        }
        
        for american, british in replacements.items():
            text = text.replace(american, british)
            text = text.replace(american.capitalize(), british.capitalize())
        
        return text
    
    def interactive_mode(self):
        """Run interactive optimizer"""
        print("\n🎓 BRITISH PROMPT OPTIMIZER")
        print("="*60)
        print("Enter ANY prompt in ANY language")
        print("I will optimize it for British education")
        print("Type 'quit' to exit\n")
        
        while True:
            user_input = input("\nOriginal prompt: ").strip()
            
            if user_input.lower() in ['quit', 'exit', 'bye']:
                print("\nCheerio! Have a brilliant day!")
                break
            
            if not user_input:
                continue
            
            print("\nOptimized prompt: ", end="", flush=True)
            optimized = self.optimize_prompt(user_input)
            print(optimized)
    
    def batch_optimize(self):
        """Optimize multiple prompts"""
        test_prompts = [
            "como fazer uma aula boa?",
            "teach math to kids",
            "explicar fotossíntese",
            "homework ideas",
            "avaliação de alunos",
            "create a science project",
            "português para crianças",
            "classroom management tips",
            "ideias para educação física",
            "how to teach Shakespeare"
        ]
        
        print("\n🔄 BATCH OPTIMIZATION")
        print("="*60)
        
        results = []
        
        for i, prompt in enumerate(test_prompts, 1):
            print(f"\n{i}. Original: {prompt}")
            optimized = self.optimize_prompt(prompt)
            print(f"   Optimized: {optimized}")
            
            results.append({
                "original": prompt,
                "optimized": optimized,
                "timestamp": datetime.now().isoformat()
            })
        
        # Save results
        output_dir = Path("prompt_optimization_results")
        output_dir.mkdir(exist_ok=True)
        
        output_file = output_dir / f"batch_optimization_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
        
        print(f"\n📁 Results saved to: {output_file}")
    
    def optimize_from_file(self, input_file, output_file):
        """Optimize prompts from JSONL file"""
        print(f"\n📂 Processing file: {input_file}")
        
        optimized_count = 0
        
        with open(input_file, 'r', encoding='utf-8') as f_in:
            with open(output_file, 'w', encoding='utf-8') as f_out:
                for line in f_in:
                    try:
                        data = json.loads(line.strip())
                        
                        # Extract user message
                        if 'messages' in data:
                            for msg in data['messages']:
                                if msg['role'] == 'user':
                                    original = msg['content']
                                    optimized = self.optimize_prompt(original)
                                    
                                    # Save optimized version
                                    output_data = {
                                        "original": original,
                                        "optimized": optimized,
                                        "timestamp": datetime.now().isoformat()
                                    }
                                    
                                    f_out.write(json.dumps(output_data, ensure_ascii=False) + '\n')
                                    optimized_count += 1
                                    
                                    if optimized_count % 100 == 0:
                                        print(f"   Processed: {optimized_count}")
                                    
                    except Exception as e:
                        logger.error(f"Error processing line: {e}")
                        continue
        
        print(f"\n✅ Optimized {optimized_count} prompts")
        print(f"📁 Saved to: {output_file}")

def main():
    optimizer = BritishPromptOptimizer()
    
    print("🎓 BRITISH PROMPT OPTIMIZER - PRODUCTION")
    print("="*60)
    print("1. Interactive mode (type prompts)")
    print("2. Batch test (10 examples)")
    print("3. Process file")
    print("4. Quick test")
    
    choice = input("\nSelect mode (1-4): ").strip()
    
    if choice == "1":
        optimizer.interactive_mode()
    elif choice == "2":
        optimizer.batch_optimize()
    elif choice == "3":
        input_file = input("Input file path: ").strip()
        output_file = input("Output file path: ").strip()
        if Path(input_file).exists():
            optimizer.optimize_from_file(input_file, output_file)
        else:
            print("File not found!")
    elif choice == "4":
        # Quick test
        test_prompts = [
            "como fazer uma aula boa?",
            "teach photosynthesis",
            "avaliação para matemática"
        ]
        print("\n🧪 QUICK TEST")
        for prompt in test_prompts:
            print(f"\nOriginal: {prompt}")
            print(f"Optimized: {optimizer.optimize_prompt(prompt)}")
    else:
        print("Invalid choice")

if __name__ == "__main__":
    main()