#!/usr/bin/env python3
"""
Filter traces from Turso/SQLite (v3) to create Modal/Synth SFT-ready .jsonl files
Supports two modes:
1. Trajectory-level filtering: Include entire trajectories above a score threshold
2. Window-based filtering: Extract high-scoring windows of actions

This is the v3 version using the new async Turso-based tracing system.
"""

import json
import argparse
import asyncio
from pathlib import Path
from typing import List, Dict, Any, Tuple, Optional
from collections import defaultdict
import numpy as np
import os
import sys
import toml
import pandas as pd

# Add synth_ai to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent.parent.parent))

from synth_ai.tracing_v3 import SessionTracer
from synth_ai.tracing_v3.turso.manager import AsyncSQLTraceManager
from synth_ai.tracing_v3.abstractions import LMCAISEvent, EnvironmentEvent, RuntimeEvent


def create_histogram(data: List[float], bins: int = 20, width: int = 60, height: int = 15, 
                    title: str = "", x_label: str = "", y_label: str = "") -> str:
    """Create a beautiful ASCII histogram."""
    if not data:
        return "No data to display"
    
    # Create histogram
    counts, edges = np.histogram(data, bins=bins)
    max_count = max(counts) if len(counts) > 0 else 1
    
    # Normalize heights
    if max_count > 0:
        heights = [int(c * height / max_count) for c in counts]
    else:
        heights = [0] * len(counts)
    
    # Build the plot
    lines = []
    
    # Title
    if title:
        lines.append(f"\n{title.center(width + 10)}")
        lines.append("=" * (width + 10))
    
    # Y-axis label
    if y_label:
        lines.append(f"{y_label}")
    
    # Plot area with y-axis
    for y in range(height, 0, -1):
        # Y-axis value
        y_val = int(max_count * y / height)
        line = f"{y_val:>6} │"
        
        # Bars
        for h in heights:
            if h >= y:
                line += "█"
            else:
                line += " "
        
        lines.append(line)
    
    # X-axis
    lines.append(f"{'':>6} └" + "─" * len(heights))
    
    # X-axis labels
    x_labels_line = " " * 8
    min_val, max_val = min(data), max(data)
    
    # Add labels at key positions
    label_positions = [0, len(heights)//4, len(heights)//2, 3*len(heights)//4, len(heights)-1]
    for i, pos in enumerate(label_positions):
        if pos < len(edges) - 1:
            val = edges[pos]
            label = f"{val:.1f}"
            # Calculate position
            target_pos = 8 + pos
            if i == 0:
                x_labels_line = label + x_labels_line[len(label):]
            elif i == len(label_positions) - 1:
                start = max(0, target_pos - len(label))
                x_labels_line = x_labels_line[:start] + label
            else:
                start = max(0, target_pos - len(label)//2)
                end = min(len(x_labels_line), start + len(label))
                if start < len(x_labels_line):
                    x_labels_line = x_labels_line[:start] + label[:end-start] + x_labels_line[end:]
    
    lines.append(x_labels_line)
    
    # X-axis label
    if x_label:
        lines.append(f"\n{x_label.center(width + 10)}")
    
    return "\n".join(lines)


def create_bar_chart(categories: List[str], values: List[int], width: int = 60, 
                     title: str = "", show_values: bool = True) -> str:
    """Create a horizontal bar chart."""
    if not categories or not values:
        return "No data to display"
    
    max_val = max(values) if values else 1
    lines = []
    
    # Title
    if title:
        lines.append(f"\n{title}")
        lines.append("=" * (width + 20))
    
    # Find longest category name for alignment
    max_cat_len = max(len(cat) for cat in categories)
    
    # Create bars
    for cat, val in zip(categories, values):
        # Normalize bar length
        bar_len = int(val * width / max_val) if max_val > 0 else 0
        bar = "█" * bar_len
        
        # Format line
        if show_values:
            line = f"{cat:<{max_cat_len}} │ {bar} {val}"
        else:
            line = f"{cat:<{max_cat_len}} │ {bar}"
        
        lines.append(line)
    
    return "\n".join(lines)


class FinetuningDataExtractorV3:
    """Extract fine-tuning data from v3 Turso traces."""
    
    def __init__(self, db_url: str):
        self.db_manager = AsyncSQLTraceManager(db_url)
        self._initialized = False
    
    async def __aenter__(self):
        await self.db_manager.initialize()
        self._initialized = True
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.db_manager.close()
    
    async def get_all_sessions(self) -> pd.DataFrame:
        """Get all session IDs from the database."""
        query = """
        SELECT DISTINCT session_id, created_at
        FROM session_traces
        ORDER BY created_at DESC
        """
        return await self.db_manager.query_traces(query)
    
    async def get_session_metrics(self, session_id: str) -> Dict[str, Any]:
        """Get metrics for a specific session."""
        # Get total reward from environment events
        # Now that rewards are properly saved in the DB, we can use them directly
        reward_query = """
        SELECT COALESCE(SUM(reward), 0) as total_reward
        FROM events
        WHERE session_id = :session_id
        AND event_type = 'environment'
        AND reward IS NOT NULL
        """
        reward_df = await self.db_manager.query_traces(reward_query, {"session_id": session_id})
        total_reward = float(reward_df['total_reward'].iloc[0]) if not reward_df.empty else 0.0
        
        # Get total tokens and cost from LM events
        lm_query = """
        SELECT 
            COALESCE(SUM(total_tokens), 0) as total_tokens,
            COALESCE(SUM(cost_usd) / 100.0, 0) as total_cost
        FROM events
        WHERE session_id = :session_id
        AND event_type = 'cais'
        """
        lm_df = await self.db_manager.query_traces(lm_query, {"session_id": session_id})
        
        total_tokens = int(lm_df['total_tokens'].iloc[0]) if not lm_df.empty else 0
        total_cost = float(lm_df['total_cost'].iloc[0]) if not lm_df.empty else 0.0
        
        return {
            'session_id': session_id,
            'total_reward': total_reward,
            'total_tokens': total_tokens,
            'total_cost': total_cost
        }
    
    async def get_session_achievements(self, session_id: str) -> List[str]:
        """Get list of achievements unlocked in a session.

        Aggregates across ALL environment events with a non-null system_state_after,
        unioning any flags that were ever true. This is more robust than inspecting
        only the last event, which can miss transient unlocks.
        """
        query = """
        SELECT system_state_after
        FROM events
        WHERE session_id = :session_id
        AND event_type = 'environment'
        AND system_state_after IS NOT NULL
        ORDER BY id ASC
        """
        df = await self.db_manager.query_traces(query, {"session_id": session_id})

        if df.empty:
            return []

        unlocked: Dict[str, bool] = {}
        for _, row in df.iterrows():
            try:
                state_after = row['system_state_after']
                if not state_after:
                    continue
                if isinstance(state_after, str):
                    state_after = json.loads(state_after)
                if not isinstance(state_after, dict):
                    continue
                public_state = state_after.get('public_state')
                if not isinstance(public_state, dict):
                    continue
                ach = public_state.get('achievements_status')
                if not isinstance(ach, dict):
                    continue
                for name, flag in ach.items():
                    if flag:
                        unlocked[name] = True
            except Exception as e:
                print(f"Error parsing achievements row: {e}")
                continue

        return [k for k, v in unlocked.items() if v]
    
    async def filter_by_achievements(self, min_achievements: int) -> List[str]:
        """Get sessions with at least min_achievements unlocked."""
        all_sessions = await self.get_all_sessions()
        qualifying_sessions = []
        
        for _, row in all_sessions.iterrows():
            session_id = row['session_id']
            achievements = await self.get_session_achievements(session_id)
            if len(achievements) >= min_achievements:
                qualifying_sessions.append(session_id)
        
        return qualifying_sessions
    
    async def extract_openai_format(self, session_ids: List[str], min_reward: float = 0.0) -> List[Dict[str, Any]]:
        """Extract training data in OpenAI format from filtered sessions."""
        training_data = []
        
        for session_id in session_ids:
            # Get messages directly from the messages table
            messages_query = """
            SELECT m.message_type, m.content, m.message_time, st.turn_number
            FROM messages m
            LEFT JOIN session_timesteps st ON m.timestep_id = st.id
            WHERE m.session_id = :session_id
            ORDER BY COALESCE(st.turn_number, m.message_time), m.id
            """
            
            messages_df = await self.db_manager.query_traces(messages_query, {"session_id": session_id})
            
            if len(messages_df) == 0:
                continue
            
            # Build conversation history
            messages = []
            system_message = None
            
            for _, row in messages_df.iterrows():
                msg_type = row['message_type']
                content = row['content']
                
                # Parse content if it's JSON (from origin_system_id format)
                try:
                    import json
                    content_data = json.loads(content)
                    if isinstance(content_data, dict) and 'payload' in content_data:
                        content = content_data['payload']
                except:
                    pass
                
                if msg_type == 'system' and system_message is None:
                    # Extract system message from the first system message
                    if isinstance(content, str):
                        system_message = content
                
                elif msg_type == 'user':
                    # Format user messages
                    if isinstance(content, dict):
                        # Convert observation dict to formatted string
                        content = self._format_observation_content(content)
                    messages.append({"role": "user", "content": str(content)})
                    
                elif msg_type == 'assistant':
                    messages.append({"role": "assistant", "content": str(content)})
            
            # Add system message at the beginning if found
            if system_message:
                messages.insert(0, {"role": "system", "content": system_message})
            
            # Only include if we have a complete conversation
            if len(messages) > 1:
                # Get total reward for this session
                reward_query = """
                SELECT COALESCE(SUM(reward), 0) as total_reward
                FROM events
                WHERE session_id = :session_id
                AND event_type = 'environment'
                AND reward IS NOT NULL
                """
                reward_df = await self.db_manager.query_traces(reward_query, {"session_id": session_id})
                total_reward = reward_df.iloc[0]['total_reward'] if len(reward_df) > 0 else 0
                
                training_data.append({
                    "messages": messages,
                    "metadata": {
                        "session_id": session_id,
                        "total_reward": float(total_reward)  # Convert to float for JSON serialization
                    }
                })
        
        return training_data

    async def extract_openai_window_format(self, session_ids: List[str]) -> List[Dict[str, Any]]:
        """Extract per-turn user→assistant pairs (window mode) for SFT.

        Emits one example per assistant message, pairing it with the preceding user
        message in the same turn (based on session_timesteps.turn_number).
        """
        window_data: List[Dict[str, Any]] = []

        for session_id in session_ids:
            messages_query = """
            SELECT st.turn_number, m.message_type, m.content, m.id AS message_id
            FROM messages m
            LEFT JOIN session_timesteps st ON m.timestep_id = st.id
            WHERE m.session_id = :session_id
            ORDER BY COALESCE(st.turn_number, m.message_time), m.id
            """
            df = await self.db_manager.query_traces(messages_query, {"session_id": session_id})
            if df is None or df.empty:
                continue

            # Parse content and group by turn_number
            parsed_rows: List[Dict[str, Any]] = []
            for _, row in df.iterrows():
                msg_type = row.get('message_type')
                content = row.get('content')
                try:
                    content_data = json.loads(content)
                    if isinstance(content_data, dict) and 'payload' in content_data:
                        content = content_data['payload']
                except Exception:
                    pass
                parsed_rows.append({
                    'turn_number': row.get('turn_number'),
                    'message_type': msg_type,
                    'content': content,
                })

            # Build windows per turn_number
            from collections import defaultdict
            turn_to_msgs: Dict[int, List[Dict[str, Any]]] = defaultdict(list)
            for r in parsed_rows:
                tn = r.get('turn_number')
                if tn is None:
                    # Skip rows that aren't associated with a turn
                    continue
                turn_to_msgs[int(tn)].append(r)

            # For each turn, find user -> assistant pair(s)
            for tn in sorted(turn_to_msgs.keys()):
                msgs = turn_to_msgs[tn]
                # find last user before first assistant
                user_content: Optional[str] = None
                assistant_content: Optional[str] = None
                for r in msgs:
                    if r['message_type'] == 'user':
                        user_content = r['content']
                    elif r['message_type'] == 'assistant' and assistant_content is None:
                        assistant_content = r['content']
                if user_content and assistant_content:
                    window_data.append({
                        'messages': [
                            { 'role': 'user', 'content': str(user_content) },
                            { 'role': 'assistant', 'content': str(assistant_content) },
                        ],
                        'metadata': {
                            'session_id': session_id,
                            'turn_number': tn,
                        }
                    })

        return window_data
    
    def _format_observation_content(self, obs: Dict[str, Any]) -> str:
        """Format observation dict into a readable string."""
        if not isinstance(obs, dict):
            return str(obs)
        
        # Extract key fields for a concise representation
        parts = []
        
        if 'inventory' in obs:
            inv = obs['inventory']
            inv_str = ", ".join([f"{k}: {v}" for k, v in inv.items() if v > 0])
            if inv_str:
                parts.append(f"Inventory: {inv_str}")
        
        if 'achievements_status' in obs:
            achievements = [k for k, v in obs['achievements_status'].items() if v]
            if achievements:
                parts.append(f"Achievements: {', '.join(achievements)}")
        
        if 'health' in obs:
            parts.append(f"Health: {obs.get('health', 0)}")
        
        return "; ".join(parts) if parts else "Empty observation"


async def filter_traces_from_turso(
    db_url: str,
    output_path: str,
    config: Dict[str, Any]
) -> Tuple[int, Dict[str, Any]]:
    """
    Filter traces from Turso/SQLite v3 database based on configuration.
    
    Returns:
        Tuple of (num_examples, statistics_dict)
    """
    mode = config.get("mode", "trajectory")
    filters = config.get("filters", {})
    
    # Extract filtering parameters
    min_reward = filters.get("min_total_reward", 0.0)
    min_achievements = filters.get("min_achievements", 0)
    max_cost = filters.get("max_cost", float('inf'))
    max_tokens = filters.get("max_tokens", float('inf'))
    
    # Modal/Synth specific: filter by model if specified
    target_models = filters.get("models", [])
    
    statistics = {
        "total_sessions": 0,
        "filtered_sessions": 0,
        "total_examples": 0,
        "reward_distribution": [],
        "token_distribution": [],
        "cost_distribution": [],
        "model_distribution": defaultdict(int)
    }
    
    async with FinetuningDataExtractorV3(db_url) as extractor:
        # Get all sessions
        all_sessions = await extractor.get_all_sessions()
        statistics["total_sessions"] = len(all_sessions)
        
        # Filter sessions based on criteria
        filtered_sessions = []
        
        for _, row in all_sessions.iterrows():
            session_id = row['session_id']
            metrics = await extractor.get_session_metrics(session_id)
            
            # Apply filters
            if metrics['total_reward'] < min_reward:
                continue
            if metrics['total_cost'] > max_cost:
                continue
            if metrics['total_tokens'] > max_tokens:
                continue
            
            # Check achievements if required
            if min_achievements > 0:
                achievements = await extractor.get_session_achievements(session_id)
                if len(achievements) < min_achievements:
                    continue
            
            # Check model filter if specified
            if target_models:
                model_query = """
                    SELECT DISTINCT model_name
                    FROM events
                    WHERE session_id = :session_id
                    AND event_type = 'cais'
                    AND model_name IS NOT NULL
                """
                model_df = await extractor.db_manager.query_traces(
                    model_query, {"session_id": session_id}
                )
                session_models = model_df['model_name'].tolist() if not model_df.empty else []
                if not any(model in target_models for model in session_models):
                    continue
            
            filtered_sessions.append(session_id)
            
            # Collect statistics
            statistics["reward_distribution"].append(metrics['total_reward'])
            statistics["token_distribution"].append(metrics['total_tokens'])
            statistics["cost_distribution"].append(metrics['total_cost'])
        
        statistics["filtered_sessions"] = len(filtered_sessions)
        
        # Extract training data
        if mode == "trajectory":
            training_data = await extractor.extract_openai_format(
                session_ids=filtered_sessions,
                min_reward=min_reward
            )
        else:  # window mode
            # For window mode, we need to implement window extraction
            # For now, use trajectory mode
            training_data = await extractor.extract_openai_format(
                session_ids=filtered_sessions,
                min_reward=min_reward
            )
        
        statistics["total_examples"] = len(training_data)
        
        # Write to output file
        output_file = Path(output_path)
        output_file.parent.mkdir(exist_ok=True)
        
        with open(output_file, 'w') as f:
            for example in training_data:
                f.write(json.dumps(example) + '\n')
        
        # Get model distribution
        model_query = """
            SELECT model_name, COUNT(*) as count
            FROM events
            WHERE event_type = 'cais'
            AND model_name IS NOT NULL
            GROUP BY model_name
        """
        model_stats = await extractor.db_manager.query_traces(model_query)
        for _, row in model_stats.iterrows():
            statistics["model_distribution"][row['model_name']] = int(row['count'])
    
    return len(training_data), statistics


def print_statistics(stats: Dict[str, Any]):
    """Print filtering statistics with visualizations."""
    print("\n" + "="*80)
    print("FILTERING STATISTICS (Modal/Synth - v3)")
    print("="*80)
    
    # Basic counts
    print(f"\nTotal sessions in database: {stats['total_sessions']}")
    print(f"Sessions after filtering: {stats['filtered_sessions']}")
    print(f"Training examples generated: {stats['total_examples']}")
    
    filter_rate = (stats['filtered_sessions'] / stats['total_sessions'] * 100) if stats['total_sessions'] > 0 else 0
    print(f"Filter pass rate: {filter_rate:.1f}%")
    
    # Reward distribution
    if stats['reward_distribution'] and any(not np.isnan(x) for x in stats['reward_distribution']):
        valid_rewards = [x for x in stats['reward_distribution'] if not np.isnan(x)]
        if valid_rewards:
            print(create_histogram(
                valid_rewards,
                bins=20,
                title="Reward Distribution",
                x_label="Total Reward",
                y_label="Count"
            ))
            
            print(f"\nReward statistics:")
            print(f"  Min: {min(valid_rewards):.2f}")
            print(f"  Max: {max(valid_rewards):.2f}")
            print(f"  Mean: {np.mean(valid_rewards):.2f}")
            print(f"  Median: {np.median(valid_rewards):.2f}")
    else:
        print("\nNo valid reward data to display.")
    
    # Token distribution
    if stats['token_distribution'] and any(not np.isnan(x) for x in stats['token_distribution']):
        valid_tokens = [x for x in stats['token_distribution'] if not np.isnan(x)]
        if valid_tokens:
            print(create_histogram(
                valid_tokens,
                bins=20,
                title="Token Usage Distribution",
                x_label="Total Tokens",
                y_label="Count"
            ))
    
    # Model distribution
    if stats['model_distribution']:
        models = list(stats['model_distribution'].keys())
        counts = list(stats['model_distribution'].values())
        print(create_bar_chart(
            models,
            counts,
            title="Model Usage",
            show_values=True
        ))
    
    print("\n" + "="*80)


def main():
    parser = argparse.ArgumentParser(
        description="Filter traces from Turso/SQLite v3 for Modal/Synth fine-tuning",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Example usage:
  # Use default config
  python filter_traces_sft_turso.py -d sqlite:///traces.db -o ft_data/training.jsonl
  
  # Use custom config file
  python filter_traces_sft_turso.py -d sqlite:///traces.db -c filter_config.toml
  
  # Override config parameters
  python filter_traces_sft_turso.py -d sqlite:///traces.db --min-reward 5.0 --max-cost 0.1
  
  # Filter by model
  python filter_traces_sft_turso.py -d sqlite:///traces.db --models "Qwen/Qwen2.5-7B-Instruct"
        """
    )
    
    parser.add_argument('-d', '--database', required=True, help='Path to Turso/SQLite database or connection URL')
    parser.add_argument('-o', '--output', default='ft_data/training_modal.jsonl', help='Output JSONL file')
    parser.add_argument('-c', '--config', help='Configuration TOML file')
    
    # Filter overrides
    parser.add_argument('--mode', choices=['trajectory', 'window'], help='Filtering mode')
    parser.add_argument('--min-reward', type=float, help='Minimum total reward')
    parser.add_argument('--min-achievements', type=int, help='Minimum achievements')
    parser.add_argument('--max-cost', type=float, help='Maximum cost')
    parser.add_argument('--max-tokens', type=int, help='Maximum tokens')
    parser.add_argument('--models', nargs='+', help='Filter by model names (e.g., Qwen/Qwen2.5-7B-Instruct)')
    
    parser.add_argument('--dry-run', action='store_true', help='Show statistics without writing output')
    
    args = parser.parse_args()
    
    # Load config
    config = {
        "mode": "trajectory",
        "filters": {
            "min_total_reward": 1.0,
            "min_achievements": 0,
            "max_cost": 10.0,
            "max_tokens": 100000,
            "models": []  # Empty means all models
        }
    }
    
    if args.config:
        with open(args.config, 'r') as f:
            loaded_config = toml.load(f)
            config.update(loaded_config)
    
    # Apply command-line overrides
    if args.mode:
        config["mode"] = args.mode
    if args.min_reward is not None:
        config["filters"]["min_total_reward"] = args.min_reward
    if args.min_achievements is not None:
        config["filters"]["min_achievements"] = args.min_achievements
    if args.max_cost is not None:
        config["filters"]["max_cost"] = args.max_cost
    if args.max_tokens is not None:
        config["filters"]["max_tokens"] = args.max_tokens
    if args.models:
        config["filters"]["models"] = args.models
    
    # Convert database path to proper URL format if needed
    db_url = args.database
    if db_url.startswith("sqlite:///"):
        # Already in URL format
        pass
    elif db_url.endswith(".db"):
        # Convert file path to URL
        db_url = f"sqlite+aiosqlite:///{db_url}"
    
    print(f"🤖 Modal/Synth Fine-Tuning Data Filter (v3)")
    print(f"Using database: {db_url}")
    print(f"Output file: {args.output}")
    print(f"Mode: {config['mode']}")
    print(f"Filters: {json.dumps(config['filters'], indent=2)}")
    
    if args.dry_run:
        print("\n🔍 DRY RUN - No output will be written")
    
    # Run filtering
    async def run():
        num_examples, stats = await filter_traces_from_turso(
            db_url,
            args.output if not args.dry_run else "/dev/null",
            config
        )
        
        # Print statistics
        print_statistics(stats)
        
        if not args.dry_run:
            print(f"\n✅ Successfully wrote {num_examples} training examples to {args.output}")
            print(f"   Ready for Modal/Synth fine-tuning!")
        else:
            print(f"\n✅ Would write {num_examples} training examples (dry run)")
    
    asyncio.run(run())


if __name__ == "__main__":
    main()