import random
import uuid
import datetime
import json
import os
from typing import Dict, List, Any, Tuple
import concurrent.futures  # Add this import for multithreading

# Import from danielson.py
from danielson import generate_single_component_evaluation
from elo import update_ratings

# Import from lowInferenceNotes.py
from lowInferenceNotes import lowInferenceNotes

# Import Argilla
import argilla as rg

# Global variables to store data temporarily
all_outputs = []  # Will be populated with outputs
current_cycle_note = None  # The current note being used for this cycle

# Initialize Argilla client
def init_argilla():
    """Initialize Argilla client"""
    try:
        print("[DEBUG] Initializing Argilla client")
        # Updated initialization method
        client = rg.Argilla(api_url=os.getenv("ARGILLA_API_URL", "http://localhost:6900"), 
                api_key=os.getenv("ARGILLA_API_KEY", "admin.apikey"))
        print(f"[DEBUG] Argilla client initialized with URL: {os.getenv('ARGILLA_API_URL', 'http://localhost:6900')}")
        
        # Test connection to verify client is working
        try:
            me = client.me
            print(f"[DEBUG] Connected to Argilla as user: {me.username}")
        except Exception as conn_error:
            print(f"[ERROR] Failed to connect to Argilla: {str(conn_error)}")
            raise conn_error
        
        
        # Create datasets if they don't exist
        try:
            # Dataset for outputs
            outputs_settings = rg.Settings(
                fields=[
                    rg.TextField(name="low_inference_notes"),
                    rg.TextField(name="output_summary"),
                    rg.TextField(name="component_id")
                ],
                questions=[
                    rg.LabelQuestion(name="quality", labels=["good", "fine", "bad"]),
                    rg.TextQuestion(name="explanation"),
                    rg.TextField(name="elo_rating", required=False)
                ],
                metadata=[
                    rg.TermsMetadataProperty(name="output_id"),
                    rg.TermsMetadataProperty(name="timestamp")
                ]
            )
            
            # Try to get dataset, create if it doesn't exist
            try:
                outputs_dataset = rg.Dataset.from_argilla("danielson_outputs")
            except:
                outputs_dataset = rg.Dataset(name="danielson_outputs", settings=outputs_settings)
                outputs_dataset.create()
                
            # Dataset for comparisons
            comparisons_settings = rg.Settings(
                fields=[
                    rg.TextField(name="low_inference_notes"),
                    rg.TextField(name="output_a"),
                    rg.TextField(name="output_b")
                ],
                questions=[
                    rg.LabelQuestion(name="preference", labels=["output_a", "output_b"]),
                    rg.TextQuestion(name="explanation")
                ],
                metadata=[
                    rg.TermsMetadataProperty(name="output_id_a"),
                    rg.TermsMetadataProperty(name="output_id_b"),
                    rg.TermsMetadataProperty(name="timestamp")
                ]
            )
            
            try:
                comparisons_dataset = rg.Dataset.from_argilla("danielson_comparisons")
            except:
                comparisons_dataset = rg.Dataset(name="danielson_comparisons", settings=comparisons_settings)
                comparisons_dataset.create()
                
        except Exception as e:
            print(f"Error creating Argilla datasets: {str(e)}")
            
    except Exception as e:
        print(f"Error initializing Argilla: {str(e)}")

def load_low_inference_notes():
    """Get the list of low inference notes"""
    return lowInferenceNotes

def generate_outputs(note, num_outputs=10):
    """Generate outputs using danielson.py for a specific note"""
    outputs = []
    
    # Component IDs for Danielson framework
    component_ids = ["1a", "1b", "1c", "1d", "1e", "1f",
                    "2a", "2b", "2c", "2d", "2e",
                    "3a", "3b", "3c", "3d", "3e"]
    
    # Select random component IDs
    selected_components = random.sample(component_ids, min(num_outputs, len(component_ids)))
    
    # Function to process a single component and return the result
    def process_component(component_id):
        try:
            # Generate output using danielson.py
            output = generate_single_component_evaluation(note, component_id)
            
            # Create output object with metadata
            output_id = str(uuid.uuid4())
            output_obj = {
                "id": output_id,
                "component_id": component_id,
                "low_inference_notes": note,
                "output": output,
                "timestamp": datetime.datetime.now().isoformat()
            }
            
            # Log to Argilla
            try:
                # Updated to use Record instead of FeedbackRecord
                record = rg.Record(
                    fields={
                        "low_inference_notes": note,
                        "output_summary": output["summary"],
                        "component_id": component_id
                    },
                    metadata={
                        "output_id": output_id,
                        "timestamp": datetime.datetime.now().isoformat()
                    }
                )
                
                # Add to Argilla dataset using the records.log() method
                client = rg.Argilla(api_url=os.getenv("ARGILLA_API_URL", "http://localhost:6900"), 
                api_key=os.getenv("ARGILLA_API_KEY", "admin.apikey"))
                dataset = client.datasets("danielson_outputs")
                dataset.records.log([record])
                
            except Exception as e:
                print(f"Error logging to Argilla: {str(e)}")
            
            return output_obj
        except Exception as e:
            print(f"Error generating output for component {component_id}: {str(e)}")
            return None
    
    # Use ThreadPoolExecutor to parallelize the process
    with concurrent.futures.ThreadPoolExecutor() as executor:
        # Submit all tasks and get future objects
        future_to_component = {executor.submit(process_component, component_id): component_id 
                              for component_id in selected_components}
        
        # Process results as they complete
        for future in concurrent.futures.as_completed(future_to_component):
            component_id = future_to_component[future]
            try:
                output_obj = future.result()
                if output_obj:
                    outputs.append(output_obj)
            except Exception as e:
                print(f"Exception for component {component_id}: {str(e)}")
    
    return outputs

def store_feedback(output_id, choice, explanation, stage):
    """Store user feedback in Argilla"""
    try:
        print(f"[DEBUG] Attempting to store feedback for output_id={output_id}, choice={choice}, stage={stage}")
        client = rg.Argilla(api_url=os.getenv("ARGILLA_API_URL", "http://localhost:6900"), 
                api_key=os.getenv("ARGILLA_API_KEY", "admin.apikey"))
        print(f"[DEBUG] Argilla client initialized: {client}")
        
        if stage == "individual":
            print(f"[DEBUG] Processing individual feedback for output_id={output_id}")
            # Get the dataset
            dataset = client.datasets("danielson_outputs")
            print(f"[DEBUG] Retrieved dataset: {dataset.name}")
            
            # Use rg.Filter for proper filtering
            filter_condition = rg.Filter(("metadata.output_id", "==", output_id))
            print(f"[DEBUG] Querying for records with filter: metadata.output_id == {output_id}")
            records_iterator = dataset.records(query=rg.Query(filter=filter_condition))
            print(f"[DEBUG] Records iterator: {records_iterator}")
            records_list = records_iterator.to_list(flatten=True)
            print(f"[DEBUG] Found {len(records_list)} records matching output_id={output_id}")
            
            if records_list:
                note = records_list[0]["low_inference_notes"]
                output = records_list[0]["output_summary"]
                component_id = records_list[0]["component_id"]
                id = records_list[0]["id"]

                # Create a record with responses we want to update
                print("Choice: ", choice)
                print("Explanation: ", explanation)
                print("Records list: ", records_list)
                print("ID: ", id)
                current_user = client.me
                print("Current user: ", current_user.id)
                
                # Create Response objects with proper format
                responses = [
                    rg.Response(question_name="quality", value=choice, user_id=current_user.id),
                    rg.Response(question_name="explanation", value=explanation, user_id=current_user.id)
                ]

                updated_record = rg.Record(
                    fields={
                        "low_inference_notes": note,
                        "output_summary": output,
                        "component_id": component_id
                    },
                    responses=responses,
                    metadata={
                        "output_id": output_id,
                        "timestamp": datetime.datetime.now().isoformat()
                    },
                    id=id
                )

                
                
                print(f"[DEBUG] Prepared update record: {updated_record}")
                
                # Log the update using the correct format
                try:
                    dataset.records.log([updated_record])
                    print(f"[DEBUG] Successfully updated record {output_id}")
                except Exception as update_error:
                    print(f"[ERROR] Failed to update record: {str(update_error)}")
                    raise update_error
                
                # Also update ELO rating
                print(f"[DEBUG] Updating ELO rating for output_id={output_id}")
                get_or_update_elo_rating(output_id, initial_rating=1500.0)
            else:
                print(f"[ERROR] No records found for output_id={output_id}")
                return False
        elif stage == "comparison":
            # For comparison feedback, we need to extract the IDs
            # Assuming choice format is "preferred_over_<other_id>"
            print(f"[DEBUG] Processing comparison feedback with choice: {choice}")
            parts = choice.split("preferred_over_")
            if len(parts) == 2:
                preferred_id = output_id
                other_id = parts[1]
                print(f"[DEBUG] Parsed preferred_id={preferred_id}, other_id={other_id}")
                
                # Create a comparison record
                # First, get the original records to extract the outputs
                dataset_outputs = client.datasets("danielson_outputs")
                print(f"[DEBUG] Retrieved outputs dataset for comparisons")
                
                # Use proper filtering for both records
                preferred_filter = rg.Filter(("metadata.output_id", "==", preferred_id))
                other_filter = rg.Filter(("metadata.output_id", "==", other_id))
                
                preferred_records_iterator = dataset_outputs.records(query=rg.Query(filter=preferred_filter))
                other_records_iterator = dataset_outputs.records(query=rg.Query(filter=other_filter))
                
                preferred_records_list = preferred_records_iterator.to_list(flatten=True)
                other_records_list = other_records_iterator.to_list(flatten=True)
                print(f"[DEBUG] Found {len(preferred_records_list)} preferred records and {len(other_records_list)} other records")
                
                if preferred_records_list and other_records_list:
                    preferred_record = preferred_records_list[0]
                    other_record = other_records_list[0]
                    current_user = client.me
                    
                    # Create comparison record with proper Response objects
                    try:
                        # Create Response objects with proper format
                        responses = [
                            rg.Response(question_name="preference", value="output_a" if preferred_id == output_id else "output_b", user_id=current_user.id),
                            rg.Response(question_name="explanation", value=explanation, user_id=current_user.id)
                        ]
                        
                        comparison_record = rg.Record(
                            fields={
                                "low_inference_notes": preferred_record["fields"]["low_inference_notes"],
                                "output_a": preferred_record["fields"]["output_summary"],
                                "output_b": other_record["fields"]["output_summary"]
                            },
                            responses=responses,
                            metadata={
                                "output_id_a": preferred_id,
                                "output_id_b": other_id,
                                "timestamp": datetime.datetime.now().isoformat()
                            }
                        )
                        print(f"[DEBUG] Created comparison record")
                        
                        # Add to comparisons dataset
                        comparisons_dataset = client.datasets("danielson_comparisons")
                        print(f"[DEBUG] Retrieved comparisons dataset")
                        comparisons_dataset.records.log([comparison_record])
                        print(f"[DEBUG] Successfully logged comparison record")
                        
                        # Update ELO ratings
                        print(f"[DEBUG] Updating ELO ratings for comparison")
                        a_wins = (preferred_id == output_id)
                        update_ratings(preferred_id, other_id, a_wins)
                        
                        return True
                    except Exception as e:
                        print(f"[ERROR] Failed to create or log comparison record: {str(e)}")
                        import traceback
                        print(f"[ERROR] Traceback: {traceback.format_exc()}")
                        return False
                else:
                    print(f"[ERROR] Could not find both records for comparison: preferred={len(preferred_records_list)}, other={len(other_records_list)}")
                    return False
            else:
                print(f"[ERROR] Invalid comparison choice format: {choice}")
                return False
        
        return True
    except Exception as e:
        print(f"Error storing feedback in Argilla: {str(e)}")
        return False

def get_or_update_elo_rating(output_id, initial_rating=None, new_rating=None):
    """Get or update ELO rating for an output"""
    try:
        print(f"[DEBUG] get_or_update_elo_rating called for output_id={output_id}, initial_rating={initial_rating}, new_rating={new_rating}")
        client = rg.Argilla(api_url=os.getenv("ARGILLA_API_URL", "http://localhost:6900"), 
                api_key=os.getenv("ARGILLA_API_KEY", "admin.apikey"))
        dataset = client.datasets("danielson_outputs")
        
        # Use proper filtering instead of string query
        filter_condition = rg.Filter(("metadata.output_id", "==", output_id))
        print(f"[DEBUG] Querying for records with filter: metadata.output_id == {output_id}")
        records_iterator = dataset.records(query=rg.Query(filter=filter_condition))
        records_list = records_iterator.to_list(flatten=True)
        print(f"[DEBUG] Found {len(records_list)} records for ELO update")
        
        if records_list:
            record = records_list[0]
            print(f"[DEBUG] Processing record ID: {record['id']}")
            
            # Get current rating if available
            current_rating = None
            if "responses" in record and record["responses"] and "elo_rating" in record["responses"]:
                current_rating = record["responses"]["elo_rating"]
            print(f"[DEBUG] Current ELO rating: {current_rating}")
            
            if new_rating is not None:
                # Create a Response object with proper format
                responses = [
                    rg.Response(question_name="elo_rating", value=str(new_rating), user_id=str(uuid.uuid4()))
                ]
                
                # Update rating with properly structured record
                elo_record = rg.Record(
                    id=record["id"],
                    responses=responses
                )
                
                print(f"[DEBUG] Updating ELO rating with record: {elo_record}")
                try:
                    dataset.records.log([elo_record])
                    print(f"[DEBUG] Successfully updated ELO rating for record {record['id']}")
                except Exception as update_error:
                    print(f"[ERROR] Failed to update ELO rating: {str(update_error)}")
                    raise update_error
                return new_rating
            elif current_rating is not None:
                # Return existing rating
                return float(current_rating)
            elif initial_rating is not None:
                # Create a Response object with proper format
                responses = [
                    rg.Response(question_name="elo_rating", value=str(initial_rating), user_id=str(uuid.uuid4()))
                ]
                
                # Set initial rating with properly structured record
                elo_record = rg.Record(
                    id=record["id"],
                    responses=responses
                )
                
                print(f"[DEBUG] Setting initial ELO rating with record: {elo_record}")
                try:
                    dataset.records.log([elo_record])
                    print(f"[DEBUG] Successfully set initial ELO rating for record {record['id']}")
                except Exception as update_error:
                    print(f"[ERROR] Failed to set initial ELO rating: {str(update_error)}")
                    raise update_error
                return initial_rating
        else:
            print(f"[ERROR] No records found for ELO update with output_id={output_id}")
        
        return None
    except Exception as e:
        print(f"[ERROR] Exception in get_or_update_elo_rating: {str(e)}")
        import traceback
        print(f"[ERROR] Traceback: {traceback.format_exc()}")
        return None

def initialize_cycle():
    """Initialize a new cycle of outputs with a new low inference note"""
    global all_outputs, current_cycle_note
    
    # Initialize Argilla if not already done
    init_argilla()
    
    # Select a random low inference note different from the current one if possible
    notes = load_low_inference_notes()
    
    if len(notes) > 1 and current_cycle_note is not None:
        # Filter out the current note to ensure we get a different one
        available_notes = [note for note in notes if note != current_cycle_note]
        current_cycle_note = random.choice(available_notes)
    else:
        # If only one note or first cycle, just choose randomly
        current_cycle_note = random.choice(notes)
    
    # Generate outputs for this note
    all_outputs = generate_outputs(current_cycle_note)
    
    return all_outputs