#!/usr/bin/env python3
"""
Task Contracts: Output Discipline for Execution Engine
Defines strict contracts for each task type to constrain LLM behavior
"""

from dataclasses import dataclass
from typing import Optional, Dict
from enum import Enum

class TaskType(Enum):
    """Supported task types with strict contracts"""
    GRAMMAR_FIX = "grammar_fix"
    MATH_CALC = "math_calc"
    CODE_GEN = "code_gen"
    EXPLANATION = "explanation"
    QUESTION_ANSWER = "qa"
    GENERAL = "general"

@dataclass
class TaskContract:
    """
    Execution contract for a task type
    
    This enforces:
    - Token budgets (prevent rambling)
    - Output style (prevent hallucination)
    - Creativity level (prevent story generation)
    - Format requirements (structured output)
    """
    task_type: TaskType
    max_tokens: int
    temperature: float  # 0.0 = deterministic, 1.0 = creative
    style: str  # "rewrite_only", "factual", "code_only", "concise"
    allow_explanation: bool
    output_format: str  # "single_line", "code_block", "paragraph"
    system_instruction: str
    
    def get_prompt_template(self, user_input: str, tool_result: Optional[str] = None) -> str:
        """Generate disciplined prompt based on contract"""
        
        if self.task_type == TaskType.GRAMMAR_FIX:
            return f"""Task: Grammar correction.
Rules:
- Output ONE corrected sentence only
- Do not add explanation
- Do not add commentary
- Preserve original meaning

Input: "{user_input}"
Corrected:"""

        elif self.task_type == TaskType.MATH_CALC:
            if tool_result:
                # Tool already gave answer - just format it
                return f"""The answer is: {tool_result}"""
            else:
                return f"""Task: Solve this math problem.
Rules:
- Show final answer only
- No explanation unless asked
- Use exact arithmetic

Problem: {user_input}
Answer:"""

        elif self.task_type == TaskType.CODE_GEN:
            return f"""Task: Write code.
Rules:
- Output code only
- Include brief comments
- No explanation outside code
- Follow best practices

Request: {user_input}
Code:
```python"""

        elif self.task_type == TaskType.EXPLANATION:
            return f"""Task: Explain concisely.
Rules:
- Maximum 3 sentences
- Focus on key points
- No unnecessary details

Question: {user_input}
Answer:"""

        elif self.task_type == TaskType.QUESTION_ANSWER:
            return f"""Task: Answer the question directly.
Rules:
- Direct answer first
- Brief explanation if needed
- No rambling

Question: {user_input}
Answer:"""

        else:  # GENERAL
            return user_input


# Task contract definitions
TASK_CONTRACTS = {
    TaskType.GRAMMAR_FIX: TaskContract(
        task_type=TaskType.GRAMMAR_FIX,
        max_tokens=48,
        temperature=0.0,
        style="rewrite_only",
        allow_explanation=False,
        output_format="single_line",
        system_instruction="You are a grammar correction tool. Output only the corrected sentence."
    ),
    
    TaskType.MATH_CALC: TaskContract(
        task_type=TaskType.MATH_CALC,
        max_tokens=32,
        temperature=0.0,
        style="factual",
        allow_explanation=False,
        output_format="single_line",
        system_instruction="You are a math solver. Output only the final answer."
    ),
    
    TaskType.CODE_GEN: TaskContract(
        task_type=TaskType.CODE_GEN,
        max_tokens=200,
        temperature=0.2,
        style="code_only",
        allow_explanation=False,
        output_format="code_block",
        system_instruction="You are a code generator. Output only clean, working code."
    ),
    
    TaskType.EXPLANATION: TaskContract(
        task_type=TaskType.EXPLANATION,
        max_tokens=128,
        temperature=0.3,
        style="concise",
        allow_explanation=True,
        output_format="paragraph",
        system_instruction="You are a concise explainer. Maximum 3 sentences."
    ),
    
    TaskType.QUESTION_ANSWER: TaskContract(
        task_type=TaskType.QUESTION_ANSWER,
        max_tokens=96,
        temperature=0.3,
        style="direct",
        allow_explanation=True,
        output_format="paragraph",
        system_instruction="You are a direct answerer. Answer first, explain briefly if needed."
    ),
    
    TaskType.GENERAL: TaskContract(
        task_type=TaskType.GENERAL,
        max_tokens=256,
        temperature=0.7,
        style="conversational",
        allow_explanation=True,
        output_format="paragraph",
        system_instruction="You are a helpful assistant."
    ),
}


def detect_task_type(query: str) -> TaskType:
    """Detect task type from query"""
    query_lower = query.lower()
    
    # Grammar fix
    if any(kw in query_lower for kw in ['fix this', 'correct this', 'grammar']):
        return TaskType.GRAMMAR_FIX
    
    # Math calculation
    if any(kw in query_lower for kw in ['calculate', 'compute', 'solve']) or \
       any(op in query for op in ['+', '-', '*', '/', '=']):
        return TaskType.MATH_CALC
    
    # Code generation
    if any(kw in query_lower for kw in ['write a', 'function', 'code', 'program', 'script']):
        return TaskType.CODE_GEN
    
    # Explanation
    if any(kw in query_lower for kw in ['explain', 'what is', 'how does']):
        return TaskType.EXPLANATION
    
    # Question answering
    if query.strip().endswith('?'):
        return TaskType.QUESTION_ANSWER
    
    return TaskType.GENERAL


def get_contract(task_type: TaskType) -> TaskContract:
    """Get contract for task type"""
    return TASK_CONTRACTS[task_type]


def should_skip_llm(task_type: TaskType, tool_result: Optional[str]) -> bool:
    """
    Determine if we should skip LLM and return tool result directly
    
    Hard-stop rule: If tool gives final answer, don't re-explain
    """
    if not tool_result:
        return False
    
    # Math: calculator gave answer, return it
    if task_type == TaskType.MATH_CALC:
        return True
    
    # Future: Add more hard-stop rules
    
    return False


def test_contracts():
    """Test task contracts"""
    print("="*70)
    print("🧪 Task Contract System Test")
    print("="*70)
    
    test_queries = [
        "Fix this: Me and him goes to store",
        "Calculate 123 * 456",
        "Write a Python function to reverse a string",
        "Explain what Python is",
        "What is 2 + 2?",
    ]
    
    for query in test_queries:
        task_type = detect_task_type(query)
        contract = get_contract(task_type)
        
        print(f"\n📝 Query: {query}")
        print(f"   Task Type: {task_type.value}")
        print(f"   Max Tokens: {contract.max_tokens}")
        print(f"   Temperature: {contract.temperature}")
        print(f"   Style: {contract.style}")
        print(f"   Format: {contract.output_format}")
        
        # Show prompt
        prompt = contract.get_prompt_template(query)
        print(f"\n   Prompt:")
        print("   " + "\n   ".join(prompt.split("\n")[:5]))
        print("   ...")


if __name__ == "__main__":
    test_contracts()