# ----------------------------------------------------------------------
# LinkedIn Message Generation Code (Refactored to mirror email structure)
# ----------------------------------------------------------------------

from typing import Dict, List, Optional
from pydantic import BaseModel
from datetime import datetime

from dhisana.schemas.sales import (
    ContentGenerationContext, 
    ConversationContext,
    CampaignContext,
    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 dhisana.utils.assistant_tool_tag import assistant_tool

# ---------------------------------------------------------------------------------------
# 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 message


# ----------------------------------------------------------------------
# LinkedIn Generation Instructions Schema (intermediate step)
# ----------------------------------------------------------------------
class LinkedInGenerationInstructions(BaseModel):
    """Instructions generated by the first step to guide LinkedIn message 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 message should sound
    key_points_to_include: str  # Main points to cover
    call_to_action: str  # What action to request


# ----------------------------------------------------------------------
# LinkedIn Connection Message Schema
# ----------------------------------------------------------------------
class LinkedInConnectMessage(BaseModel):
    body: str

# ----------------------------------------------------------------------
# Cleanup function (similar to cleanup_email_context)
# ----------------------------------------------------------------------
def cleanup_linkedin_context(linkedin_context: ContentGenerationContext) -> ContentGenerationContext:
    """
    Return a copy of ContentGenerationContext without unneeded or sensitive fields 
    for LinkedIn generation.
    """
    clone_context = linkedin_context.copy(deep=True)
    
    # Example: remove irrelevant external data for this context
    if clone_context.external_known_data:
        clone_context.external_known_data.external_openai_vector_store_id = None

    # Remove extra fields on the lead_info if desired
    clone_context.lead_info.task_ids = None
    clone_context.lead_info.email_validation_status = None
    clone_context.lead_info.linkedin_validation_status = None
    clone_context.lead_info.research_status = None
    clone_context.lead_info.enchrichment_status = None

    return clone_context

# ----------------------------------------------------------------------
# Known Framework Variations (fallback if user instructions are not provided)
# ----------------------------------------------------------------------
LINKEDIN_FRAMEWORK_VARIATIONS = [
    "Use a friendly intro mentioning their role and a quick reason to connect.",
    "Use social proof: reference an industry success story or insight.",
    "Use a brief mention of mutual interest or connection.",
    "Use P-S-B style (Pain, Solution, Benefit) but under 40 words.",
    "Use a 3-Bullet Approach: Industry/Pain, Value, Simple Ask."
]

# ----------------------------------------------------------------------
# Step 1: Generate LinkedIn instructions using gpt-4.1 (larger context)
# ----------------------------------------------------------------------
async def generate_linkedin_instructions(
    linkedin_context: ContentGenerationContext,
    variation_text: str,
    tool_config: Optional[List[Dict]] = None,
) -> LinkedInGenerationInstructions:
    """
    Step 1 of 2-step LinkedIn message generation process.
    Uses gpt-4.1 (larger context length) to process all context and generate
    concise instructions for message generation.
    """
    cleaned_context = cleanup_linkedin_context(linkedin_context)

    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()

    # Construct the prompt for generating instructions
    instructions_prompt = f"""
Generate LinkedIn connection message instructions based on the following business context.
The final message must be under 40 words.

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 []}

MESSAGING FRAMEWORK:
{variation_text}

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)
- sender_context: Sender name and relevant info for signature
- campaign_context: Product or service, value proposition, key pain points
- conversation_context: Summary of any prior conversation if exists
- tone_and_style: Professional, friendly, etc.
- key_points_to_include: Main points for the message (keep brief for 40 word limit)
- call_to_action: What action the recipient should take
    """

    vector_store_id = (
        linkedin_context.external_known_data.external_openai_vector_store_id
        if linkedin_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=LinkedInGenerationInstructions,
            vector_store_id=vector_store_id,
            model=DEFAULT_OPENAI_MODEL,
            tool_config=tool_config,
            use_cache=linkedin_context.message_instructions.use_cache if linkedin_context.message_instructions else True
        )
    else:
        instructions_response, instructions_status = await get_structured_output_internal(
            prompt=instructions_prompt,
            response_format=LinkedInGenerationInstructions,
            model=DEFAULT_OPENAI_MODEL,
            tool_config=tool_config,
            use_cache=linkedin_context.message_instructions.use_cache if linkedin_context.message_instructions else True
        )

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


# ----------------------------------------------------------------------
# Step 2: Generate LinkedIn message using gpt-5.1-chat (better content)
# ----------------------------------------------------------------------
async def generate_linkedin_from_instructions(
    instructions: LinkedInGenerationInstructions,
    tool_config: Optional[List[Dict]] = None,
    use_cache: bool = True,
) -> LinkedInConnectMessage:
    """
    Step 2 of 2-step LinkedIn message generation process.
    Uses gpt-5.1-chat (better content generation) to generate the message
    from the condensed instructions.
    """
    message_prompt = f"""
Generate a personalized LinkedIn connection request message based on these specifications.
The message must be under 40 words total.

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}

OUTPUT REQUIREMENTS:
- Output must be JSON with "body" field only.
- The entire message body must be under 40 words total.
- Plain text only, no HTML tags.
- Use salutation Hi followed by first name, unless specified otherwise.
- Use only the information provided.
- Do not use em dash in the output.
    """

    # Generate message using gpt-5.1-chat (better content generation)
    message_response, message_status = await get_structured_output_internal(
        prompt=message_prompt,
        response_format=LinkedInConnectMessage,
        model=DEFAULT_EMAIL_GEN_MODEL,
        tool_config=tool_config,
        use_cache=use_cache
    )

    if message_status != "SUCCESS":
        raise Exception("Error: Could not generate LinkedIn message in step 2.")
    
    return message_response


# ----------------------------------------------------------------------
# Core function to generate a LinkedIn copy (2-step process)
# ----------------------------------------------------------------------
async def generate_personalized_linkedin_copy(
    linkedin_context: ContentGenerationContext,
    variation_text: str,
    tool_config: Optional[List[Dict]] = None,
) -> dict:
    """
    Generate a personalized LinkedIn connection message using a 2-step process:

    Step 1: Use gpt-4.1 (larger context length) to process all context and generate
            concise instructions for message generation.
    
    Step 2: Use gpt-5.1-chat (better content generation) to generate the actual
            LinkedIn message from the condensed instructions.
    """
    cleaned_context = cleanup_linkedin_context(linkedin_context)
    lead_data = cleaned_context.lead_info or Lead()
    sender_data = cleaned_context.sender_info or SenderInfo()

    # Step 1: Generate instructions using gpt-4.1
    instructions = await generate_linkedin_instructions(
        linkedin_context=linkedin_context,
        variation_text=variation_text,
        tool_config=tool_config,
    )

    # Step 2: Generate message using gpt-5.1-chat
    use_cache = linkedin_context.message_instructions.use_cache if linkedin_context.message_instructions else True
    message_response = await generate_linkedin_from_instructions(
        instructions=instructions,
        tool_config=tool_config,
        use_cache=use_cache,
    )

    # Wrap in MessageItem for consistency
    response_item = MessageItem(
        message_id="",  # fill in if you have an ID
        thread_id="",
        sender_name=sender_data.sender_full_name or "",
        sender_email=sender_data.sender_email or "",
        receiver_name=lead_data.full_name or "",
        receiver_email=lead_data.email or "",
        iso_datetime=datetime.utcnow().isoformat(),
        subject="Hi",
        body=message_response.body
    )

    return response_item.model_dump()

# ----------------------------------------------------------------------
# Primary function to generate multiple LinkedIn message variations
# (similar to generate_personalized_email)
# ----------------------------------------------------------------------
@assistant_tool
async def generate_personalized_linkedin_message(
    linkedin_context: ContentGenerationContext,
    number_of_variations: int = 3,
    tool_config: Optional[List[Dict]] = None
) -> List[dict]:
    """
    Generate multiple variations of a personalized LinkedIn connection message 
    using the provided context and instructions.

    :param linkedin_context: Consolidated context for LinkedIn generation
    :param number_of_variations: Number of variations to produce
    :param tool_config: Optional config for tool or vector store
    :return: A list of dictionaries, each containing 'subject' and 'body' only
    """
    message_instructions = linkedin_context.message_instructions or MessageGenerationInstructions()
    user_instructions_exist = bool((message_instructions.instructions_to_generate_message or "").strip())

    linkedin_variations = []
    
    for i in range(number_of_variations):
        # If user provided custom instructions, use them for each variation
        # Otherwise, pick from fallback frameworks in a round-robin fashion
        if user_instructions_exist:
            variation_text = message_instructions.instructions_to_generate_message
        else:
            variation_text = LINKEDIN_FRAMEWORK_VARIATIONS[i % len(LINKEDIN_FRAMEWORK_VARIATIONS)]
        
        personalized_copy = await generate_personalized_linkedin_copy(
            linkedin_context=linkedin_context,
            variation_text=variation_text,
            tool_config=tool_config
        )
        linkedin_variations.append(personalized_copy)

    return linkedin_variations
