"""
Refinement Agent for iterative email improvements.
Handles follow-up prompts to modify and refine generated emails.
"""
import json
from typing import Dict, Any
from langchain_core.messages import HumanMessage, SystemMessage
from agents.base_agent import BaseAgent


class RefinementAgent(BaseAgent):
    """Agent responsible for refining emails based on user feedback."""
    
    def __init__(self):
        super().__init__(temperature=0.7)
    
    @property
    def name(self) -> str:
        return "refinement_agent"
    
    def _get_system_prompt(self) -> str:
        return """You are an expert email refinement assistant.
Your task is to modify an existing email based on user instructions.

You can handle various refinement requests such as:
- Making the email shorter or longer
- Changing the tone (more formal, casual, assertive)
- Adding or removing specific content
- Improving clarity or flow
- Fixing grammar or spelling
- Changing the greeting or closing
- Adding urgency or softening the message
- Making it more polite or more direct

IMPORTANT: Preserve the overall structure and intent of the email while making the requested changes.

Respond ONLY with a valid JSON object:
{
    "subject": "updated subject line",
    "greeting": "updated greeting",
    "body": "updated body content",
    "closing": "updated closing",
    "signature": "updated signature",
    "changes_made": ["list of changes that were made"]
}"""

    def refine(self, current_email: str, user_instruction: str, 
               current_tone: str = "formal") -> Dict[str, Any]:
        """Refine the email based on user instruction."""
        
        context = f"""CURRENT EMAIL:
{current_email}

CURRENT TONE: {current_tone}

USER INSTRUCTION: {user_instruction}

Please modify the email according to the user's instruction while maintaining coherence."""

        messages = [
            SystemMessage(content=self._get_system_prompt()),
            HumanMessage(content=context)
        ]
        
        try:
            response = self.llm.invoke(messages)
            content = response.content.strip()
            
            # Clean up response
            if content.startswith("```json"):
                content = content[7:]
            if content.startswith("```"):
                content = content[3:]
            if content.endswith("```"):
                content = content[:-3]
            content = content.strip()
            
            parsed_data = json.loads(content)
            
            # Build full email
            full_email_parts = []
            if parsed_data.get("subject"):
                full_email_parts.append(f"Subject: {parsed_data['subject']}\n")
            full_email_parts.append(parsed_data.get("greeting", "Hello,"))
            full_email_parts.append("")
            full_email_parts.append(parsed_data.get("body", ""))
            full_email_parts.append("")
            full_email_parts.append(parsed_data.get("closing", "Best regards,"))
            if parsed_data.get("signature"):
                full_email_parts.append(parsed_data["signature"])
            
            full_email = "\n".join(full_email_parts)
            
            return {
                "success": True,
                "email": full_email,
                "subject": parsed_data.get("subject", ""),
                "changes_made": parsed_data.get("changes_made", [])
            }
            
        except json.JSONDecodeError as e:
            return {
                "success": False,
                "error": f"Failed to parse refinement response: {str(e)}",
                "email": current_email
            }
        except Exception as e:
            return {
                "success": False,
                "error": f"Refinement error: {str(e)}",
                "email": current_email
            }
    
    def process(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Process method for compatibility with base agent."""
        return state


# Quick action handlers
QUICK_ACTIONS = {
    "shorter": "Make the email significantly shorter and more concise while keeping the key points.",
    "longer": "Expand the email with more details, context, and elaboration.",
    "more_formal": "Make the email more formal and professional in tone.",
    "more_casual": "Make the email more casual and friendly in tone.",
    "more_assertive": "Make the email more direct, confident, and action-oriented.",
    "more_polite": "Make the email more polite and courteous.",
    "add_urgency": "Add a sense of urgency to the email without being pushy.",
    "soften": "Soften the tone and make it less demanding or aggressive.",
    "fix_grammar": "Fix any grammar, spelling, or punctuation errors.",
    "improve_clarity": "Improve the clarity and readability of the email.",
    "add_cta": "Add a clear call-to-action at the end of the email.",
    "remove_fluff": "Remove unnecessary words and phrases to make it more direct."
}


def get_quick_action_instruction(action_key: str) -> str:
    """Get the instruction for a quick action."""
    return QUICK_ACTIONS.get(action_key, action_key)