import json
import os
import random
import time
import requests
from typing import List, Dict, Any
import logging
from PyPDF2 import PdfReader

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("dataset_generation.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# DeepSeek API configuration
DEEPSEEK_API_URL = "https://api.deepseek.com/v1/chat/completions"
DEEPSEEK_MODEL = "deepseek-chat"

def load_json(path):
    """Load a JSON file."""
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

def extract_text_from_pdf(pdf_path):
    """Extract text from a PDF file."""
    logger.info(f"Extracting text from {pdf_path}...")
    reader = PdfReader(pdf_path)
    text = ""
    for i, page in enumerate(reader.pages):
        content = page.extract_text()
        if content:
            text += content + "\n"
        if (i + 1) % 10 == 0:
            logger.info(f"Processed {i + 1} pages...")
    logger.info(f"Finished extracting text from {len(reader.pages)} pages.")
    return text

def clean_text(text: str) -> str:
    """Clean text by removing special characters and normalizing whitespace."""
    import re
    import unicodedata
    
    # Remove null bytes and control characters
    text = re.sub(r'[\x00-\x1F\x7F-\x9F]', '', text)
    
    # Normalize quotes
    text = text.replace('"', '"').replace('"', '"')
    text = text.replace(''', "'").replace(''', "'")
    
    # Fix spacing around punctuation
    text = re.sub(r'\s+([.,!?])', r'\1', text)
    text = re.sub(r'([.,!?])\s*([A-Z])', r'\1 \2', text)
    
    # Normalize whitespace
    text = re.sub(r'\s+', ' ', text)
    
    # Fix common formatting issues
    text = text.replace(' .', '.')
    text = text.replace(' ,', ',')
    text = text.replace(' !', '!')
    text = text.replace(' ?', '?')
    
    # Remove any remaining special characters
    text = ''.join(char for char in text if unicodedata.category(char)[0] != 'C')
    
    return text.strip()

def clean_dataset_item(item: Dict[str, Any]) -> Dict[str, Any]:
    """Clean a dataset item by cleaning all string fields."""
    if 'prompt' in item:
        item['prompt'] = clean_text(item['prompt'])
    if 'completion' in item:
        item['completion'] = clean_text(item['completion'])
    for key, value in item.items():
        if isinstance(value, str):
            item[key] = clean_text(value)
    return item

def call_deepseek_api(prompt: str, api_key: str) -> str:
    """Call the DeepSeek API with the given prompt."""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    
    data = {
        "model": DEEPSEEK_MODEL,
        "messages": [
            {"role": "system", "content": "You are an expert in creating high-quality training data for fine-tuning language models. You write in a clear, professional style."},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.7,
        "max_tokens": 2000
    }
    
    try:
        response = requests.post(DEEPSEEK_API_URL, headers=headers, json=data)
        response.raise_for_status()
        result = response.json()
        return result["choices"][0]["message"]["content"]
    except Exception as e:
        logger.error(f"Error calling DeepSeek API: {e}")
        return None

def extract_json_from_response(response: str) -> List[Dict]:
    """Extract JSON array from the API response."""
    try:
        # First try to parse the entire response as JSON
        try:
            return json.loads(response)
        except json.JSONDecodeError:
            pass
        
        # Try to find a JSON array in the response
        import re
        json_match = re.search(r'\[\s*\{.*\}\s*\]', response, re.DOTALL)
        if json_match:
            json_str = json_match.group(0)
            return json.loads(json_str)
        
        # If no JSON array found, try to extract individual examples
        examples = []
        pattern = r'\{[\s"]*"prompt"[\s"]*:[\s"]*"([^"]*)"[\s"]*,[\s"]*"completion"[\s"]*:[\s"]*"([^"]*)"[\s"]*\}'
        matches = re.findall(pattern, response, re.DOTALL)
        
        for prompt, completion in matches:
            examples.append({
                "prompt": prompt,
                "completion": completion
            })
        
        if examples:
            return examples
        
        logger.error("Could not extract JSON from response")
        return []
    except Exception as e:
        logger.error(f"Error extracting JSON from response: {e}")
        return []

def generate_book_content_examples(book_text: str, brand_guidelines: Dict, tone_keywords: Dict, num_examples: int, api_key: str) -> List[Dict]:
    """Generate examples based on book content."""
    logger.info(f"Generating {num_examples} examples based on book content...")
    
    # Extract key concepts and themes from the book
    system_prompt = f"""
    You are an expert in creating training data for fine-tuning language models.
    I need you to generate examples that capture the writing style and expertise of Adriana James.
    
    Brand Guidelines:
    - Tone: {brand_guidelines['tone']}
    - Voice: {brand_guidelines['voice']}
    - Brand Values: {', '.join(brand_guidelines['brand_values'])}
    - Key Phrases: {', '.join(brand_guidelines['key_phrases'])}
    - Avoided Phrases: {', '.join(brand_guidelines['avoided_phrases'])}
    
    Here is a summary of the book content:
    {book_text[:2000]}...
    
    Generate {num_examples} training examples in the following format:
    {{
        "prompt": "A simple user request without mentioning style (e.g., 'Write a motivational social media post')",
        "completion": "A response that matches Adriana James' writing style, tone, and expertise"
    }}
    
    IMPORTANT: The prompts should be simple user requests WITHOUT mentioning style or Adriana James.
    The completions should ALWAYS be in Adriana James' style, regardless of the prompt.
    
    The completions should:
    1. Match the tone and voice specified in the brand guidelines
    2. Incorporate the brand values and key phrases
    3. Avoid the specified phrases
    4. Reflect the expertise and insights from the book
    5. Be realistic and professional
    
    Return the examples as a JSON array.
    """
    
    response = call_deepseek_api(system_prompt, api_key)
    if not response:
        logger.error("Failed to generate book content examples")
        return []
    
    examples = extract_json_from_response(response)
    if examples:
        logger.info(f"Generated {len(examples)} book content examples")
    else:
        logger.error("Failed to parse book content examples")
    
    return examples

def generate_campaign_examples(campaigns: List[Dict], brand_guidelines: Dict, tone_keywords: Dict, num_examples: int, api_key: str) -> List[Dict]:
    """Generate examples based on past campaigns."""
    logger.info(f"Generating {num_examples} examples based on past campaigns...")
    
    # Format campaigns for the prompt
    campaigns_text = json.dumps(campaigns, indent=2)
    
    system_prompt = f"""
    You are an expert in creating training data for fine-tuning language models.
    I need you to generate examples that capture the marketing style of Adriana James.
    
    Brand Guidelines:
    - Tone: {brand_guidelines['tone']}
    - Voice: {brand_guidelines['voice']}
    - Brand Values: {', '.join(brand_guidelines['brand_values'])}
    - Key Phrases: {', '.join(brand_guidelines['key_phrases'])}
    - Avoided Phrases: {', '.join(brand_guidelines['avoided_phrases'])}
    
    Here are examples of past campaigns:
    {campaigns_text}
    
    Generate {num_examples} training examples in the following format:
    {{
        "prompt": "A simple user request without mentioning style (e.g., 'Write an email about our new leadership program')",
        "completion": "A marketing piece that matches Adriana James' writing style and tone"
    }}
    
    IMPORTANT: The prompts should be simple user requests WITHOUT mentioning style or Adriana James.
    The completions should ALWAYS be in Adriana James' style, regardless of the prompt.
    
    The completions should:
    1. Match the tone and voice specified in the brand guidelines
    2. Incorporate the brand values and key phrases
    3. Avoid the specified phrases
    4. Be similar in style to the past campaigns
    5. Be realistic and professional
    
    Return the examples as a JSON array.
    """
    
    response = call_deepseek_api(system_prompt, api_key)
    if not response:
        logger.error("Failed to generate campaign examples")
        return []
    
    examples = extract_json_from_response(response)
    if examples:
        logger.info(f"Generated {len(examples)} campaign examples")
    else:
        logger.error("Failed to parse campaign examples")
    
    return examples

def generate_style_examples(brand_guidelines: Dict, tone_keywords: Dict, num_examples: int, api_key: str) -> List[Dict]:
    """Generate examples focused on style alignment."""
    logger.info(f"Generating {num_examples} style alignment examples...")
    
    system_prompt = f"""
    You are an expert in creating training data for fine-tuning language models.
    I need you to generate examples that help a model learn to match Adriana James' writing style.
    
    Brand Guidelines:
    - Tone: {brand_guidelines['tone']}
    - Voice: {brand_guidelines['voice']}
    - Brand Values: {', '.join(brand_guidelines['brand_values'])}
    - Key Phrases: {', '.join(brand_guidelines['key_phrases'])}
    - Avoided Phrases: {', '.join(brand_guidelines['avoided_phrases'])}
    
    Generate {num_examples} training examples in the following format:
    {{
        "prompt": "A simple user request without mentioning style (e.g., 'Write a product description for our coaching program')",
        "completion": "The content in Adriana James' style"
    }}
    
    IMPORTANT: The prompts should be simple user requests WITHOUT mentioning style or Adriana James.
    The completions should ALWAYS be in Adriana James' style, regardless of the prompt.
    
    The examples should include:
    1. Content that would normally be written in different tones, but is rewritten in Adriana James' tone
    2. Content that would normally be written in different voices, but is rewritten in Adriana James' voice
    3. Content that incorporates brand values and key phrases
    4. Content that avoids specified phrases
    5. Content that sounds professional and expert
    
    Return the examples as a JSON array.
    """
    
    response = call_deepseek_api(system_prompt, api_key)
    if not response:
        logger.error("Failed to generate style examples")
        return []
    
    examples = extract_json_from_response(response)
    if examples:
        logger.info(f"Generated {len(examples)} style examples")
    else:
        logger.error("Failed to parse style examples")
    
    return examples

def generate_synthetic_examples(brand_guidelines: Dict, tone_keywords: Dict, num_examples: int, api_key: str) -> List[Dict]:
    """Generate synthetic examples for various content types."""
    logger.info(f"Generating {num_examples} synthetic examples...")
    
    system_prompt = f"""
    You are an expert in creating training data for fine-tuning language models.
    I need you to generate examples that help a model learn to create various types of content in Adriana James' style.
    
    Brand Guidelines:
    - Tone: {brand_guidelines['tone']}
    - Voice: {brand_guidelines['voice']}
    - Brand Values: {', '.join(brand_guidelines['brand_values'])}
    - Key Phrases: {', '.join(brand_guidelines['key_phrases'])}
    - Avoided Phrases: {', '.join(brand_guidelines['avoided_phrases'])}
    
    Generate {num_examples} training examples in the following format:
    {{
        "prompt": "A simple user request without mentioning style (e.g., 'Write a social media post about our new leadership program')",
        "completion": "The generated content in Adriana James' style"
    }}
    
    IMPORTANT: The prompts should be simple user requests WITHOUT mentioning style or Adriana James.
    The completions should ALWAYS be in Adriana James' style, regardless of the prompt.
    
    Include a variety of content types:
    - Email subject lines
    - Email bodies
    - Social media posts
    - Blog titles
    - Blog introductions
    - Product descriptions
    
    And a variety of topics:
    - Leadership
    - Professional development
    - Career advancement
    - Business strategy
    - Personal branding
    - Work-life balance
    - Decision making
    - Communication skills
    
    And a variety of audiences:
    - Professional women
    - Business leaders
    - Entrepreneurs
    - Executives
    - Career-driven individuals
    - Industry professionals
    
    The completions should:
    1. Match the tone and voice specified in the brand guidelines
    2. Incorporate the brand values and key phrases
    3. Avoid the specified phrases
    4. Be realistic and professional
    
    Return the examples as a JSON array.
    """
    
    response = call_deepseek_api(system_prompt, api_key)
    if not response:
        logger.error("Failed to generate synthetic examples")
        return []
    
    examples = extract_json_from_response(response)
    if examples:
        logger.info(f"Generated {len(examples)} synthetic examples")
    else:
        logger.error("Failed to parse synthetic examples")
    
    return examples

def save_dataset(dataset: List[Dict], filename: str):
    """Save a dataset to a JSON file."""
    logger.info(f"Saving dataset to {filename}...")
    
    # Clean the dataset
    cleaned_dataset = [clean_dataset_item(item) for item in dataset]
    
    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    
    # Save to file
    with open(filename, 'w', encoding='utf-8') as f:
        json.dump(cleaned_dataset, f, indent=2, ensure_ascii=False)
    
    logger.info(f"Dataset saved with {len(cleaned_dataset)} items.")

def main():
    # Get API key from environment variable
    api_key = "sk-776f96462d954d4f98c4ee1a493ac9ad"
    if not api_key:
        logger.error("DEEPSEEK_API_KEY environment variable not set")
        return
    
    # Create output directory
    os.makedirs("datasets", exist_ok=True)
    
    # Load brand guidelines and tone keywords
    brand_guidelines = load_json("data/style_guidelines/brand_guidelines.json")
    tone_keywords = load_json("data/style_guidelines/tone_keywords.json")
    
    # Extract text from the book
    book_text = extract_text_from_pdf("data/book.pdf")
    
    # Load past campaigns
    campaigns = load_json("data/past_campaigns/sample_campaign.json")
    
    # Generate examples for each category
    book_examples = generate_book_content_examples(book_text, brand_guidelines, tone_keywords, 30, api_key)
    campaign_examples = generate_campaign_examples(campaigns, brand_guidelines, tone_keywords, 20, api_key)
    style_examples = generate_style_examples(brand_guidelines, tone_keywords, 30, api_key)
    synthetic_examples = generate_synthetic_examples(brand_guidelines, tone_keywords, 40, api_key)
    
    # Create fallback examples if API calls failed
    if not book_examples:
        logger.info("Creating fallback book content examples...")
        book_examples = create_fallback_book_examples(book_text, brand_guidelines, tone_keywords, 30)
    
    if not campaign_examples:
        logger.info("Creating fallback campaign examples...")
        campaign_examples = create_fallback_campaign_examples(campaigns, brand_guidelines, tone_keywords, 20)
    
    if not style_examples:
        logger.info("Creating fallback style examples...")
        style_examples = create_fallback_style_examples(brand_guidelines, tone_keywords, 30)
    
    if not synthetic_examples:
        logger.info("Creating fallback synthetic examples...")
        synthetic_examples = create_fallback_synthetic_examples(brand_guidelines, tone_keywords, 40)
    
    # Save individual datasets
    save_dataset(book_examples, "datasets/stage1_book_content.json")
    save_dataset(campaign_examples, "datasets/stage2_campaign_content.json")
    save_dataset(synthetic_examples, "datasets/stage3_synthetic_content.json")
    save_dataset(style_examples, "datasets/stage4_style_alignment.json")
    
    # Combine all datasets
    combined_dataset = book_examples + campaign_examples + synthetic_examples + style_examples
    
    # Shuffle the combined dataset
    random.shuffle(combined_dataset)
    
    # Save the combined dataset
    save_dataset(combined_dataset, "datasets/adriana_finetune_dataset.json")
    
    logger.info("Dataset generation complete!")

def create_fallback_book_examples(book_text: str, brand_guidelines: Dict, tone_keywords: Dict, num_examples: int) -> List[Dict]:
    """Create fallback book content examples if API call fails."""
    logger.info(f"Creating {num_examples} fallback book content examples...")
    
    # Extract paragraphs from the book
    paragraphs = [p.strip() for p in book_text.split('\n\n') if p.strip()]
    
    # Create realistic prompts
    realistic_prompts = [
        "Write a motivational social media post.",
        "Draft a personal growth blog intro.",
        "Create a coaching-oriented email body.",
        "Write an inspirational paragraph for professionals.",
        "Craft a message about values and transformation.",
        "Write a personal development newsletter paragraph."
    ]
    
    # Create examples
    examples = []
    for _ in range(min(num_examples, len(paragraphs))):
        prompt = random.choice(realistic_prompts)
        completion = random.choice(paragraphs)
        examples.append({
            "prompt": prompt,
            "completion": completion
        })
    
    logger.info(f"Created {len(examples)} fallback book content examples")
    return examples

def create_fallback_campaign_examples(campaigns: List[Dict], brand_guidelines: Dict, tone_keywords: Dict, num_examples: int) -> List[Dict]:
    """Create fallback campaign examples if API call fails."""
    logger.info(f"Creating {num_examples} fallback campaign examples...")
    
    if not isinstance(campaigns, list):
        campaigns = [campaigns]
    
    examples = []
    for campaign in campaigns:
        prompt = (
            f"Write a {campaign['type'].replace('_', ' ')} for {campaign['target_audience']}.\n"
            f"Title: {campaign['title']}\n"
            f"Focus: {brand_guidelines['tone']} and {brand_guidelines['voice']}.\n"
        )
        
        if campaign['type'] == 'email_campaign':
            completion = (
                f"Subject: {campaign['content']['subject_line']}\n\n"
                f"{campaign['content']['body']}\n\n"
                f"Call to Action: {campaign['content']['call_to_action']}"
            )
        elif campaign['type'] == 'social_media':
            completion = campaign['content'].get('post', '')
        else:
            completion = json.dumps(campaign['content'], indent=2)
        
        examples.append({
            "prompt": prompt,
            "completion": completion
        })
    
    logger.info(f"Created {len(examples)} fallback campaign examples")
    return examples

def create_fallback_style_examples(brand_guidelines: Dict, tone_keywords: Dict, num_examples: int) -> List[Dict]:
    """Create fallback style examples if API call fails."""
    logger.info(f"Creating {num_examples} fallback style examples...")
    
    examples = []
    for _ in range(num_examples // 3):
        examples.append({
            "prompt": "Rewrite the following content to match a confident and professional tone:\nThis product is the best ever! It's cheap and guaranteed to make you rich!",
            "completion": "This program is designed to deliver measurable results and meaningful growth for professionals. It's a strategic investment in your future success."
        })
        examples.append({
            "prompt": "Rewrite the following to sound more like expert guidance:\nYou could maybe try this if you want.",
            "completion": "This is a proven approach recommended for professionals who want to achieve clear, measurable results."
        })
        examples.append({
            "prompt": "Rewrite the following to avoid over-promising language:\nThis is the number one product that will solve everything instantly!",
            "completion": "This well-designed solution is backed by years of experience and offers a practical framework for long-term success."
        })
    
    logger.info(f"Created {len(examples)} fallback style examples")
    return examples

def create_fallback_synthetic_examples(brand_guidelines: Dict, tone_keywords: Dict, num_examples: int) -> List[Dict]:
    """Create fallback synthetic examples if API call fails."""
    logger.info(f"Creating {num_examples} fallback synthetic examples...")
    
    content_types = ["email_subject", "email_body", "social_media_post", "blog_title", "blog_intro", "product_description"]
    target_audiences = ["professional women", "business leaders", "entrepreneurs", "executives", "career-driven individuals"]
    topics = ["leadership", "professional development", "career advancement", "business strategy", "personal branding"]
    
    templates = {
        "social_media_post": [
            "Write a social media post about {topic}.",
            "Create a post for {audience} on {topic}.",
        ],
        "email_body": [
            "Write an email body for {audience} on {topic}.",
        ],
        "blog_intro": [
            "Write a blog intro about {topic}.",
        ],
        "product_description": [
            "Write a product description focused on {topic} for {audience}.",
        ],
        "email_subject": [
            "Write an email subject line about {topic}.",
        ],
        "blog_title": [
            "Write a blog title about {topic}.",
        ]
    }
    
    examples = []
    for _ in range(num_examples):
        content_type = random.choice(content_types)
        audience = random.choice(target_audiences)
        topic = random.choice(topics)
        prompt = random.choice(templates[content_type]).format(topic=topic, audience=audience)
        
        # Create a simple completion based on brand guidelines
        completion = f"This {content_type.replace('_', ' ')} about {topic} for {audience} reflects our commitment to {', '.join(brand_guidelines['brand_values'][:2])}."
        
        examples.append({
            "prompt": prompt,
            "completion": completion
        })
    
    logger.info(f"Created {len(examples)} fallback synthetic examples")
    return examples

if __name__ == "__main__":
    main() 