"""
Example Query Library for Natural Language Query System.

Contains curated examples of natural language queries mapped to OuthadEdge code.
Used for few-shot learning and RAG retrieval.
"""

from typing import List, Dict


class ExampleQueryLibrary:
    """
    Library of example NL queries → code pairs for few-shot learning
    """

    @staticmethod
    def get_examples() -> List[Dict[str, str]]:
        """
        Get all example queries with code and expected answers.

        Returns:
            list: Example queries with query, code, and answer
        """
        return [
            {
                'query': 'What is the overall conversion rate?',
                'code': '''# Calculate conversion rate
df = stream.to_dataframe()
total_users = df['user_id'].nunique()
converted_users = df[df['event'] == 'payment_done']['user_id'].nunique()
conversion_rate = (converted_users / total_users) * 100
result = "{:.2f}%".format(conversion_rate)
''',
                'answer': '17.4%',
                'description': 'Calculate overall conversion rate from all users to payment_done event',
                'tags': ['conversion', 'metrics', 'basic']
            },
            {
                'query': 'Show funnel from homepage to purchase',
                'code': '''# Build conversion funnel
funnel = stream.funnel()
funnel.fit(stages=['homepage', 'product_view', 'add_to_cart', 'checkout', 'purchase'])
result = funnel.values
visualization = funnel  # Returns visualization object
''',
                'answer': 'Funnel created with conversion rates at each stage',
                'description': 'Create multi-stage conversion funnel visualization',
                'tags': ['funnel', 'conversion', 'visualization']
            },
            {
                'query': 'What is the retention rate for weekly cohorts?',
                'code': '''# Calculate weekly cohort retention
cohorts = stream.cohorts(
    cohort_start_unit='W',
    cohort_period=(1, 'W')
)
result = cohorts.cohort_matrix
visualization = cohorts  # Returns heatmap
''',
                'answer': 'Cohort retention matrix showing weekly retention rates',
                'description': 'Analyze user retention using weekly cohorts',
                'tags': ['cohorts', 'retention', 'visualization']
            },
            {
                'query': 'How many users from mobile vs desktop?',
                'code': '''# Count users by device type
df = stream.to_dataframe()

# Assume we have device column, else need to filter by user_agent
# For this example, we'll use event patterns
mobile_users = df[df['event'].str.contains('mobile', case=False, na=False)]['user_id'].nunique()
total_users = df['user_id'].nunique()
desktop_users = total_users - mobile_users

result = f"Mobile: {mobile_users}, Desktop: {desktop_users}"
''',
                'answer': 'Mobile: 4,523, Desktop: 8,231',
                'description': 'Segment users by device type',
                'tags': ['segmentation', 'device', 'users']
            },
            {
                'query': 'What events lead to highest conversion?',
                'code': '''# Find events correlated with conversion
df = stream.to_dataframe()

# Users who converted
converted_users = df[df['event'] == 'purchase']['user_id'].unique()

# Events before conversion
pre_conversion_events = df[df['user_id'].isin(converted_users)]
event_counts = pre_conversion_events.groupby('event')['user_id'].nunique().sort_values(ascending=False)

result = event_counts.head(10)
''',
                'answer': 'Top events: product_review (85%), add_to_cart (78%), compare (65%)',
                'description': 'Identify events with highest correlation to conversion',
                'tags': ['conversion', 'analysis', 'correlation']
            },
            {
                'query': 'Segment users by behavior',
                'code': '''# Behavioral segmentation using clustering
clusters = stream.clusters()
X = clusters.extract_features(method='tfidf', ngram_range=(1,2))
clusters.fit(method='kmeans', n_clusters=4, X=X, random_state=42)

# Get cluster assignments
segments = clusters.cluster_mapping

result = f"Created {len(segments['cluster_id'].unique())} behavioral segments"
visualization = clusters  # Returns cluster visualization
''',
                'answer': '4 segments: Power Users, Browsers, Deal Seekers, At-Risk',
                'description': 'Create behavioral segments using ML clustering',
                'tags': ['clustering', 'segmentation', 'ml', 'visualization']
            },
            {
                'query': 'What is the average time to purchase?',
                'code': '''# Calculate average time from first event to purchase
df = stream.to_dataframe()

# Get first and purchase events per user
user_times = df.groupby('user_id').agg({
    'timestamp': ['min', 'max']
}).reset_index()

user_times.columns = ['user_id', 'first_event', 'last_event']

# Filter users who made purchase
purchase_users = df[df['event'] == 'purchase']['user_id'].unique()
converted_times = user_times[user_times['user_id'].isin(purchase_users)]

# Calculate time difference
converted_times['time_to_purchase'] = (converted_times['last_event'] - converted_times['first_event']).dt.total_seconds() / 3600

avg_hours = converted_times['time_to_purchase'].mean()
result = f"Average time to purchase: {avg_hours:.1f} hours"
''',
                'answer': 'Average time to purchase: 12.3 hours',
                'description': 'Calculate average time from first interaction to conversion',
                'tags': ['metrics', 'time', 'conversion']
            },
            {
                'query': 'Show transition graph of user flow',
                'code': '''# Create transition graph visualization
graph = stream.transition_graph()
result = graph
''',
                'answer': 'Interactive transition graph showing user flow between events',
                'description': 'Visualize user journey as network graph',
                'tags': ['visualization', 'transition', 'flow']
            },
            {
                'query': 'What is the churn rate?',
                'code': '''# Calculate churn rate
df = stream.to_dataframe()

# Define churn as no activity in last 30 days
from datetime import datetime, timedelta
cutoff_date = df['timestamp'].max() - timedelta(days=30)

recent_users = df[df['timestamp'] > cutoff_date]['user_id'].nunique()
total_users = df['user_id'].nunique()

churn_rate = ((total_users - recent_users) / total_users) * 100
result = f"Churn rate: {churn_rate:.2f}%"
''',
                'answer': 'Churn rate: 23.5%',
                'description': 'Calculate percentage of users who stopped engaging',
                'tags': ['churn', 'retention', 'metrics']
            },
            {
                'query': 'Compare conversion rates by segment',
                'code': '''# A/B test comparison
stats = stream.stattests()
result = stats.chi2_test(
    groups=['segment_a', 'segment_b'],
    event='purchase'
)
''',
                'answer': 'Chi-square test results showing statistical significance',
                'description': 'Statistical comparison of conversion across segments',
                'tags': ['ab-test', 'statistics', 'comparison']
            },
            {
                'query': 'Get monthly active users',
                'code': '''# Calculate MAU
df = stream.to_dataframe()
df['month'] = df['timestamp'].dt.to_period('M')

mau_by_month = df.groupby('month')['user_id'].nunique()
result = mau_by_month
''',
                'answer': 'Monthly active users by month',
                'description': 'Count unique active users per month',
                'tags': ['metrics', 'mau', 'engagement']
            },
            {
                'query': 'What are the most common user paths?',
                'code': '''# Analyze common sequences
sequences = stream.sequences(ngram_range=(2, 3), threshold=['count', 100])
result = sequences
''',
                'answer': 'Top user paths with frequency counts',
                'description': 'Identify most frequent event sequences',
                'tags': ['sequences', 'paths', 'patterns']
            },
        ]

    @staticmethod
    def get_by_tag(tag: str) -> List[Dict[str, str]]:
        """
        Get examples filtered by tag.

        Args:
            tag: Tag to filter by (e.g., 'conversion', 'funnel')

        Returns:
            list: Filtered examples
        """
        all_examples = ExampleQueryLibrary.get_examples()
        return [
            ex for ex in all_examples
            if tag.lower() in [t.lower() for t in ex.get('tags', [])]
        ]

    @staticmethod
    def get_by_category(category: str) -> List[Dict[str, str]]:
        """
        Get examples by category.

        Args:
            category: Category name (metrics, visualization, segmentation, etc.)

        Returns:
            list: Examples in that category
        """
        category_map = {
            'metrics': ['conversion', 'metrics', 'mau', 'churn'],
            'visualization': ['visualization', 'funnel', 'transition'],
            'segmentation': ['segmentation', 'clustering', 'cohorts'],
            'analysis': ['analysis', 'correlation', 'statistics'],
        }

        tags = category_map.get(category.lower(), [])
        all_examples = ExampleQueryLibrary.get_examples()

        return [
            ex for ex in all_examples
            if any(tag in ex.get('tags', []) for tag in tags)
        ]

    @staticmethod
    def search(query_text: str) -> List[Dict[str, str]]:
        """
        Search examples by query text similarity.

        Args:
            query_text: Search text

        Returns:
            list: Matching examples
        """
        all_examples = ExampleQueryLibrary.get_examples()
        query_lower = query_text.lower()

        # Simple keyword matching
        matches = []
        for ex in all_examples:
            if (
                query_lower in ex['query'].lower()
                or query_lower in ex['description'].lower()
                or any(query_lower in tag for tag in ex.get('tags', []))
            ):
                matches.append(ex)

        return matches


def get_all_tags() -> List[str]:
    """Get all unique tags from examples."""
    all_tags = set()
    for example in ExampleQueryLibrary.get_examples():
        all_tags.update(example.get('tags', []))
    return sorted(list(all_tags))
