# Import necessary modules
import html as html_lib
import re
from typing import Dict, List, Optional
from pydantic import BaseModel

from dhisana.schemas.sales import CampaignContext, ContentGenerationContext, ConversationContext, Lead, MessageGenerationInstructions, MessageItem, SenderInfo
from dhisana.utils.generate_structured_output_internal import (
    get_structured_output_internal,
    get_structured_output_with_assistant_and_vector_store
)
from datetime import datetime
from pydantic import BaseModel, ConfigDict

# ---------------------------------------------------------------------------------------
# CONSTANTS
# ---------------------------------------------------------------------------------------
DEFAULT_OPENAI_MODEL = "gpt-4.1"  # Larger context length - used for generating instructions
DEFAULT_EMAIL_GEN_MODEL = "gpt-5.1-chat"  # Better content generation - used for generating email

# -----------------------------------------------------------------------------
# Email Generation Instructions schema (intermediate step)
# -----------------------------------------------------------------------------
class EmailGenerationInstructions(BaseModel):
    """Instructions generated by the first step to guide email generation."""
    recipient_context: str  # Key info about the lead/recipient
    sender_context: str  # Key info about the sender
    campaign_context: str  # Product, value prop, pain points, etc.
    conversation_context: str  # Summary of any existing conversation
    tone_and_style: str  # How the email should sound
    key_points_to_include: str  # Main points to cover in the email
    call_to_action: str  # What action to request
    formatting_requirements: str  # HTML vs plain text, signature format, etc.
    
    model_config = ConfigDict(extra="forbid")


# -----------------------------------------------------------------------------
# Email Copy schema
# -----------------------------------------------------------------------------
class EmailCopy(BaseModel):
    subject: str
    body: str
    body_html: Optional[str] = None
 
    model_config = ConfigDict(extra="forbid")


def _html_to_plain_text(html_content: str) -> str:
    """Simple HTML to text conversion to backfill plain body."""
    if not html_content:
        return ""
    # Remove tags and normalize whitespace
    text = re.sub(r"<[^>]+>", " ", html_content)
    text = html_lib.unescape(text)
    # Collapse repeated whitespace/newlines
    lines = [line.strip() for line in text.splitlines()]
    return "\n".join([line for line in lines if line])

# -----------------------------------------------------------------------------
# Utility to Clean Up Context (if needed)
# -----------------------------------------------------------------------------
def cleanup_email_context(email_context: ContentGenerationContext) -> ContentGenerationContext:
    """
    Return a copy of ContentGenerationContext without sensitive or irrelevant fields.
    Modify or remove fields as necessary for your project.
    """
    clone_context = email_context.copy(deep=True)
    # For demonstration, no sensitive fields in new classes by default.
    # Adjust if you want to remove certain fields (like 'sender_bio', etc.).
    if clone_context.external_known_data:
        clone_context.external_known_data.external_openai_vector_store_id = None
    return clone_context

# -----------------------------------------------------------------------------
# Known Framework Variations (fallback if user instructions are not provided)
# -----------------------------------------------------------------------------
FRAMEWORK_VARIATIONS = [
    "Use PAS (Problem, Agitate, Solve) framework to write up the email.",
    "Use VETO framework (Value, Evidence, Tie, Offer). Highlight how our product addresses the company's current goals.",
    "Use AIDA framework (Attention, Interest, Desire, Action) to compose the email.",
    "Use SPIN (Situation, Problem, Implication, Need-Payoff) framework to write the email.",
    "Use BANT (Budget, Authority, Need, Timeline) framework to write the email.",
    "Use P-S-B (Pain, Solution, Benefit) with a 3-Bullet Approach. Keep it concise.",
    "Use Hook, Insight, Offer framework to write the email."
]

# -----------------------------------------------------------------------------
# Step 1: Generate email instructions using gpt-4.1 (larger context)
# -----------------------------------------------------------------------------
async def generate_email_instructions(
    email_context: ContentGenerationContext,
    message_instructions: MessageGenerationInstructions,
    variation_text: str,
    tool_config: Optional[List[Dict]] = None,
) -> EmailGenerationInstructions:
    """
    Step 1 of 2-step email generation process.
    Uses gpt-4.1 (larger context length) to process all context and generate
    concise instructions for email generation.
    """
    cleaned_context = cleanup_email_context(email_context)

    # Check if user provided custom instructions
    user_custom_instructions = (message_instructions.instructions_to_generate_message or "").strip()
    use_custom_instructions = bool(user_custom_instructions)

    # Decide final instructions: user-provided or fallback variation
    if use_custom_instructions:
        selected_instructions = user_custom_instructions
    else:
        selected_instructions = variation_text

    # Pull out fields or fallback to empty if None
    lead_data = cleaned_context.lead_info or Lead()
    sender_data = cleaned_context.sender_info or SenderInfo()
    campaign_data = cleaned_context.campaign_context or CampaignContext()
    conversation_data = cleaned_context.current_conversation_context or ConversationContext()

    html_note = (
        f"\n        HTML Template to follow: {message_instructions.html_template}"
        if getattr(message_instructions, "html_template", None)
        else ""
    )
    
    allow_html = getattr(message_instructions, "allow_html", False)
    format_requirement = "HTML email with body_html field" if allow_html else "Plain text only (no HTML tags)"

    # Construct the prompt for generating instructions
    instructions_prompt = f"""
Generate email writing instructions based on the following business context.

LEAD INFORMATION:
{lead_data.dict()}

SENDER INFORMATION:
Full Name: {sender_data.sender_full_name or ''}
First Name: {sender_data.sender_first_name or ''}
Last Name: {sender_data.sender_last_name or ''}
Bio: {sender_data.sender_bio or ''}
Appointment Booking URL: {sender_data.sender_appointment_booking_url or ''}

CAMPAIGN INFORMATION:
Product Name: {campaign_data.product_name or ''}
Value Proposition: {campaign_data.value_prop or ''}
Call To Action: {campaign_data.call_to_action or ''}
Pain Points: {campaign_data.pain_points or []}
Proof Points: {campaign_data.proof_points or []}
Email Triage Guidelines: {campaign_data.email_triage_guidelines or ''}
LinkedIn Triage Guidelines: {campaign_data.linkedin_triage_guidelines or ''}

MESSAGING FRAMEWORK:
{selected_instructions}{html_note}

CONVERSATION HISTORY:
Email Thread: {conversation_data.current_email_thread or ''}
LinkedIn Thread: {conversation_data.current_linkedin_thread or ''}

OUTPUT FIELDS:
- recipient_context: Key details about the lead (name, company, role, relevant background)
- sender_context: Sender name and relevant info for signature
- campaign_context: Product or service being promoted, value proposition, key pain points
- conversation_context: Summary of any prior conversation if exists
- tone_and_style: Professional, friendly, etc.
- key_points_to_include: Specific points to mention in the email
- call_to_action: What action the recipient should take
- formatting_requirements: Format is {format_requirement}. Salutation format, signature requirements.
    """

    # Check if a vector store is available
    vector_store_id = (email_context.external_known_data.external_openai_vector_store_id
                       if email_context.external_known_data else None)

    instructions_response = None
    instructions_status = ""

    # Generate instructions using gpt-4.1 (larger context)
    if vector_store_id:
        instructions_response, instructions_status = await get_structured_output_with_assistant_and_vector_store(
            prompt=instructions_prompt,
            response_format=EmailGenerationInstructions,
            vector_store_id=vector_store_id,
            model=DEFAULT_OPENAI_MODEL,
            tool_config=tool_config,
            use_cache=email_context.message_instructions.use_cache if email_context.message_instructions else True
        )
    else:
        instructions_response, instructions_status = await get_structured_output_internal(
            prompt=instructions_prompt,
            response_format=EmailGenerationInstructions,
            model=DEFAULT_OPENAI_MODEL,
            tool_config=tool_config,
            use_cache=email_context.message_instructions.use_cache if email_context.message_instructions else True
        )

    if instructions_status != "SUCCESS":
        raise Exception("Error: Could not generate email instructions in step 1.")
    
    return instructions_response


# -----------------------------------------------------------------------------
# Step 2: Generate email using gpt-5.1-chat (better content generation)
# -----------------------------------------------------------------------------
async def generate_email_from_instructions(
    instructions: EmailGenerationInstructions,
    message_instructions: MessageGenerationInstructions,
    tool_config: Optional[List[Dict]] = None,
    use_cache: bool = True,
) -> EmailCopy:
    """
    Step 2 of 2-step email generation process.
    Uses gpt-5.1-chat (better content generation) to generate the email
    from the condensed instructions.
    """
    allow_html = getattr(message_instructions, "allow_html", False)
    
    if allow_html:
        output_requirements = """
        OUTPUT REQUIREMENTS:
        - Output must be JSON with "subject", "body", and "body_html" fields.
        - "body_html" should be clean HTML suitable for email (no external assets), inline styles welcome.
        - "body" must be the plain-text equivalent of "body_html".
        """
    else:
        output_requirements = """
        OUTPUT REQUIREMENTS:
        - Output must be JSON with "subject" and "body" fields only.
        - In the subject or body DO NOT include any HTML tags like <a>, <b>, <i>, etc.
        - The body and subject should be in plain text.
        - If there is a link provided, use it as is without wrapping in any HTML tags.
        """

    html_template_note = ""
    if allow_html and getattr(message_instructions, "html_template", None):
        html_template_note = f"\nHTML Template to follow: {message_instructions.html_template}"

    email_prompt = f"""
Generate a personalized business email based on these specifications:

RECIPIENT CONTEXT:
{instructions.recipient_context}

SENDER CONTEXT:
{instructions.sender_context}

CAMPAIGN CONTEXT:
{instructions.campaign_context}

CONVERSATION CONTEXT:
{instructions.conversation_context}

TONE AND STYLE:
{instructions.tone_and_style}

KEY POINTS TO INCLUDE:
{instructions.key_points_to_include}

CALL TO ACTION:
{instructions.call_to_action}

FORMATTING REQUIREMENTS:
{instructions.formatting_requirements}

{html_template_note}

{output_requirements}

ADDITIONAL REQUIREMENTS:
- Keep it concise and relevant. No placeholders.
- Do not include internal references or content identifiers in the email.
- Use only the information provided.
- Ensure the body text is well-formatted with proper newlines.
- Do not use em dash in the output.
    """

    # Generate email using gpt-5.1-chat (better content generation)
    email_response, email_status = await get_structured_output_internal(
        prompt=email_prompt,
        response_format=EmailCopy,
        model=DEFAULT_EMAIL_GEN_MODEL,
        tool_config=tool_config,
        use_cache=use_cache
    )

    if email_status != "SUCCESS":
        raise Exception("Error: Could not generate email in step 2.")
    
    return email_response


# -----------------------------------------------------------------------------
# Core function to generate an email copy (2-step process)
# -----------------------------------------------------------------------------
async def generate_personalized_email_copy(
    email_context: ContentGenerationContext,
    message_instructions: MessageGenerationInstructions,
    variation_text: str,
    tool_config: Optional[List[Dict]] = None,
) -> dict:
    """
    Generate a personalized email using a 2-step process:

    Step 1: Use gpt-4.1 (larger context length) to process all context and generate
            concise instructions for email generation.
    
    Step 2: Use gpt-5.1-chat (better content generation) to generate the actual
            email from the condensed instructions.

    This approach leverages gpt-4.1's larger context window to process extensive
    lead/campaign data, and gpt-5.1-chat's superior content generation for the
    final email output.
    """
    # Step 1: Generate instructions using gpt-4.1
    instructions = await generate_email_instructions(
        email_context=email_context,
        message_instructions=message_instructions,
        variation_text=variation_text,
        tool_config=tool_config,
    )

    # Step 2: Generate email using gpt-5.1-chat
    use_cache = email_context.message_instructions.use_cache if email_context.message_instructions else True
    email_response = await generate_email_from_instructions(
        instructions=instructions,
        message_instructions=message_instructions,
        tool_config=tool_config,
        use_cache=use_cache,
    )

    plain_body = email_response.body
    html_body = getattr(email_response, "body_html", None)
    if not plain_body and html_body:
        plain_body = _html_to_plain_text(html_body)

    response_item = MessageItem(
        message_id="",  # or some real ID if you have it
        thread_id="",
        sender_name=email_context.sender_info.sender_full_name or "",
        sender_email=email_context.sender_info.sender_email or "",
        receiver_name=email_context.lead_info.full_name or "",
        receiver_email=email_context.lead_info.email or "",
        iso_datetime=datetime.utcnow().isoformat(),
        subject=email_response.subject,
        body=plain_body,
        html_body=html_body if getattr(message_instructions, "allow_html", False) else None,
    )
    return response_item.model_dump()

# -----------------------------------------------------------------------------
# Primary function to generate multiple variations
# -----------------------------------------------------------------------------
async def generate_personalized_email(
    generation_context: ContentGenerationContext,
    number_of_variations: int = 3,
    tool_config: Optional[List[Dict]] = None
) -> List[dict]:
    """
    Generates multiple personalized email variations based on the given context and instructions.

    Parameters:
        - email_context: The consolidated context for email generation.
        - message_instructions: User-supplied instructions or templates for generating the message.
        - number_of_variations: How many email variations to produce.
        - tool_config: Optional tool configuration for the LLM calls.

    Returns:
        A list of dictionaries, each containing:
          {
            "subject": "string",
            "body": "string"
          }
    """
    email_variations = []
    message_instructions = generation_context.message_instructions 
    user_instructions_exist = bool(
        (message_instructions.instructions_to_generate_message or "").strip()
    )

    for i in range(number_of_variations):
        try:
            # If user provided instructions, use them for each variation
            # (skip the internal frameworks).
            if user_instructions_exist:
                variation_text = message_instructions.instructions_to_generate_message or ""
            else:
                # Otherwise, pick from known frameworks (circular indexing)
                variation_text = FRAMEWORK_VARIATIONS[i % len(FRAMEWORK_VARIATIONS)]

            email_copy = await generate_personalized_email_copy(
                email_context=generation_context,
                message_instructions=message_instructions,
                variation_text=variation_text,
                tool_config=tool_config
            )
            email_variations.append(email_copy)

        except Exception as e:
            raise e
    return email_variations
