"""
Flask-based web application for visualizing training logs.
"""

from flask import Flask, jsonify, request, send_from_directory
from pathlib import Path
import json
import os
import re
from typing import List, Dict, Any
from collections import defaultdict
import pandas as pd
import numpy as np

# Get the directory where app.py is located
app_dir = Path(__file__).parent
# Path to the built frontend
frontend_dist = app_dir / 'dist'

app = Flask(__name__, static_folder=str(frontend_dist), static_url_path='')

# Default storage path
DEFAULT_STORAGE_PATH = Path.cwd() / ".investigation"


def get_storage_path():
    """Get the storage path for training logs."""
    # Check environment variable first
    env_storage = os.environ.get('Investigation_STORAGE')
    if env_storage:
        return Path(env_storage)

    storage_env = Path.cwd() / ".investigation"
    if storage_env.exists():
        return storage_env
    return DEFAULT_STORAGE_PATH


def aggregate_parameters(storage_path: Path) -> Dict[str, Dict[str, set]]:
    """
    Aggregate all unique parameter values from JSON files in commit directories.
    Looks for files matching pattern [A-Z0-9]{8}.json in each commit directory.

    Returns:
        Dictionary mapping commit_hash to parameter aggregations
        Format: {commit_hash: {param_name: set(values)}}
    """
    aggregated_params = {}

    if not storage_path.exists():
        return aggregated_params

    # Iterate through commit directories
    for commit_dir in storage_path.iterdir():
        if not commit_dir.is_dir():
            continue

        commit_hash = commit_dir.name
        param_aggregation = defaultdict(set)

        # Find all JSON files matching the pattern [A-Z0-9]{8}.json
        for json_file in commit_dir.glob("[A-Z0-9][A-Z0-9][A-Z0-9][A-Z0-9][A-Z0-9][A-Z0-9][A-Z0-9][A-Z0-9].json"):
            try:
                with open(json_file, 'r', encoding='utf-8') as f:
                    job_configs = json.load(f)

                    # Each JSON file contains multiple job configurations
                    for job_hash, job_params in job_configs.items():
                        for param_name, param_value in job_params.items():
                            # Convert value to string for consistency in filtering
                            # if (param_name in ["num_cpus", "num_gpus", "time"]):
                            #     continue
                            param_aggregation[param_name].add(str(param_value))
            except (json.JSONDecodeError, IOError):
                continue

        # Convert sets to sorted lists for JSON serialization
        aggregated_params[commit_hash] = {param: sorted(list(values)) for param, values in param_aggregation.items()}

    return aggregated_params


def get_job_parameters(storage_path: Path, job_name: str, commit_hash: str) -> Dict[str, Any]:
    """
    Get parameters for a specific job from JSON files.

    Args:
        storage_path: Storage directory path
        job_name: Job hash (e.g., "435W8AYN")
        commit_hash: Commit hash

    Returns:
        Dictionary of job parameters or empty dict if not found
    """
    commit_dir = storage_path / commit_hash
    if not commit_dir.exists():
        return {}

    # Search through all JSON files in the commit directory
    for json_file in commit_dir.glob("[A-Z0-9][A-Z0-9][A-Z0-9][A-Z0-9][A-Z0-9][A-Z0-9][A-Z0-9][A-Z0-9].json"):
        try:
            with open(json_file, 'r', encoding='utf-8') as f:
                job_configs: Dict[str, Dict[str, Any]] = json.load(f)
                if job_name in job_configs:
                    job_config = job_configs[job_name]
                    job_config.pop("num_cpus", None)
                    job_config.pop("num_gpus", None)
                    job_config.pop("time", None)
                    job_config.pop("seed", None)
                    job_config.pop("requeue", None)
                    job_config.pop("compile_model", None)
                    return job_config
        except (json.JSONDecodeError, IOError):
            continue

    return {}


def find_csv_files(storage_path: Path) -> List[Dict[str, Any]]:
    """
    Find all CSV files in the storage directory.
    Specifically looks for files in .investigation/commithash/storage/*.csv structure.

    Returns:
        List of dictionaries with file information including job parameters
    """
    csv_files = []
    if not storage_path.exists():
        return csv_files

    # Look for CSV files in commit_hash/storage subdirectories
    for commit_dir in storage_path.iterdir():
        if not commit_dir.is_dir():
            continue

        storage_dir = commit_dir / "storage"
        if not storage_dir.exists():
            continue

        instances = []
        for csv_file in storage_dir.glob("*.csv"):
            job_name = csv_file.stem
            commit_hash = commit_dir.name
            try:
                if (csv_file.stat().st_size == 0):
                    print(f"Skipping empty CSV: {csv_file}")
                    continue
            except Exception as e:
                # print(f"Error accessing file {csv_file}: {e}")
                continue
            instances.append((csv_file, job_name, commit_hash))
        for csv_file in storage_dir.glob("*/*.csv"):
            inst_dir = csv_file.parent
            job_name = inst_dir.stem
            commit_hash = commit_dir.name
            try:
                if (csv_file.stat().st_size == 0):
                    print(f"Skipping empty CSV: {csv_file}")
                    continue
            except Exception as e:
                # print(f"Error accessing file {csv_file}: {e}")
                continue
            instances.append((csv_file, job_name, commit_hash))
        for csv_file, job_name, commit_hash in instances:
            # Get job parameters from JSON files
            job_params = get_job_parameters(storage_path, job_name, commit_hash)

            csv_files.append({
                "path": str(csv_file.relative_to(storage_path)),
                "full_path": str(csv_file),
                "name": job_name,
                "parent": csv_file.parent.name,
                "commit": commit_hash,
                "parameters": job_params
            })

    return csv_files


def read_csv_data(file_path: Path, raw: bool, columns: List[str]) -> Dict[str, Any]:
    """
    Read CSV file and return structured data. Uses pandas for performance.

    Args:
        file_path: Full path to CSV file
        raw: If True, return raw data without sampling; otherwise return sampled data

    Returns:
        Dictionary with headers and data rows
    """
    try:
        # It's recommended to add 'import pandas as pd' at the top of the file.

        # Validate file path to prevent path traversal
        file_path_obj = file_path.resolve()
        storage_path = get_storage_path().resolve()

        # Ensure the file is within the storage directory
        if not str(file_path_obj).startswith(str(storage_path)):
            return {"error": "Access denied: file outside storage directory", "headers": [], "data": []}

        # Use pandas to read the CSV, which is much faster
        df = pd.read_csv(file_path_obj)

        if (len(columns) > 0):
            # Filter to only requested columns plus global_step
            required_cols = ["global_step"] + [col for col in columns if col in df.columns]
            df = df[required_cols]

        # remove rows where all columns except global_step are NaN
        print(f"Initial rows: {len(df)}, raw={raw}, columns={columns}")
        df = df.dropna(how='all', subset=[col for col in df.columns if col != "global_step"], axis=0)
        print(f"Rows after dropping all-NaN: {len(df)}, raw={raw}, columns={columns}")
        # Replace NaN with None for JSON compatibility
        df = df.replace({np.nan: None})

        total_rows = len(df)

        if raw:
            # Return raw data without sampling
            # Limit to 50000 rows to prevent memory issues with very large files
            max_raw_rows = 50000
            if total_rows > max_raw_rows:
                df = df.head(max_raw_rows)
            headers = df.columns.tolist()
            data = df.to_dict(orient='records')
            return {"headers": headers, "data": data, "row_count": total_rows, "truncated": total_rows > max_raw_rows}

        df_index = df.set_index("global_step")

        def sample_column(col: str):
            target_col = df_index[col].dropna()
            total_points = len(target_col)
            max_points = 500
            if total_points > max_points:
                indices = np.linspace(0, total_points - 1, num=max_points, dtype=int)
                dt = target_col.iloc[indices]
            else:
                dt = target_col
            # print(f"Column {col}: total points={total_points}, sampled points={len(dt)}")
            # print(dt)
            return dt

        sampled_data = pd.concat([sample_column(col) for col in df_index.columns],
                                 axis=1).reset_index().replace({np.nan: None})
        sampled_data.sort_values(by="global_step", inplace=True)
        headers = sampled_data.columns.tolist()
        # print(sampled_data)
        data = sampled_data.to_dict(orient='records')

        return {"headers": headers, "data": data, "row_count": total_rows}
    except FileNotFoundError:
        return {"error": "File not found", "headers": [], "data": []}
    except Exception as e:
        # Catch pandas errors (e.g., EmptyDataError) and other exceptions
        return {"error": f"Failed to read or parse CSV file: {e}", "headers": [], "data": []}


def aggregate_data(instances: List[str], column: str, method: str = "mean") -> Dict[str, Any]:
    """
    Aggregate data across multiple instances.

    Args:
        instances: List of file paths to aggregate
        column: Column name to aggregate
        method: Aggregation method (mean, min, max)

    Returns:
        Aggregated data
    """
    all_data = defaultdict(list)

    for instance_path in instances:
        csv_data = read_csv_data(instance_path)
        if "error" in csv_data:
            continue

        for row in csv_data["data"]:
            if column in row and isinstance(row[column], (int, float)):
                # Use global_step or index as x-axis
                x_value = row.get("global_step", row.get("step", len(all_data[column])))
                all_data[column].append((x_value, row[column]))

    # Sort by x-axis value
    if column in all_data:
        all_data[column].sort(key=lambda x: x[0])

    # Perform aggregation
    # Group by x-value first
    grouped = defaultdict(list)
    for x, y in all_data.get(column, []):
        grouped[x].append(y)

    # Apply aggregation function
    if method == "mean":
        result = [(x, sum(values) / len(values) if len(values) > 0 else 0) for x, values in sorted(grouped.items())]
    elif method == "min":
        result = [(x, min(values) if len(values) > 0 else 0) for x, values in sorted(grouped.items())]
    elif method == "max":
        result = [(x, max(values) if len(values) > 0 else 0) for x, values in sorted(grouped.items())]
    else:
        result = all_data.get(column, [])

    return {"column": column, "method": method, "data": result}


@app.route('/')
def index():
    """Render the main visualization page."""
    return send_from_directory(app.static_folder, 'index.html')


@app.route('/<path:path>')
def serve_static(path):
    """Serve static files from the Vue build."""
    if path and (frontend_dist / path).exists():
        return send_from_directory(app.static_folder, path)
    # For Vue router or non-existent files, serve index.html
    return send_from_directory(app.static_folder, 'index.html')


@app.route('/api/instances')
def list_instances():
    """List all available training instances."""
    storage_path = get_storage_path()
    csv_files = find_csv_files(storage_path)
    return jsonify({"instances": csv_files, "count": len(csv_files)})


@app.route('/api/data/<path:instance_path>')
def get_instance_data(instance_path: str):
    """Get data for a specific instance.

    Query parameters:
        raw: If 'true', return raw data without sampling
    """
    storage_path = get_storage_path()
    raw = request.args.get('raw', 'false').lower() == 'true'
    cols = request.args.getlist('columns')
    # print(f"Requested columns: {cols}")

    # Validate path to prevent directory traversal
    try:
        full_path = (storage_path / instance_path).resolve()
        storage_path_resolved = storage_path.resolve()

        # Ensure the file is within the storage directory
        if not str(full_path).startswith(str(storage_path_resolved)):
            return jsonify({"error": "Access denied"}), 403

        if not full_path.exists():
            return jsonify({"error": "File not found"}), 404

        csv_data = read_csv_data(full_path, raw=raw, columns=cols)

        # Check if there was an error in reading
        if "error" in csv_data and csv_data["error"]:
            return jsonify({"error": "Failed to read file"}), 500

        return jsonify(csv_data)
    except Exception:
        return jsonify({"error": "Invalid path"}), 400


@app.route('/api/aggregate', methods=['POST'])
def aggregate_instances():
    """Aggregate data across multiple instances."""
    data = request.json
    instances = data.get('instances', [])
    column = data.get('column', 'episodic_return')
    method = data.get('method', 'mean')

    storage_path = get_storage_path()
    full_paths = [str(storage_path / inst) for inst in instances]

    aggregated = aggregate_data(full_paths, column, method)
    return jsonify(aggregated)


@app.route('/api/filter', methods=['POST'])
def filter_instances():
    """Filter instances based on criteria including parameters."""
    data = request.json
    pattern = data.get('pattern', '')
    param_filters = data.get('param_filters', {})  # {param_name: [value1, value2, ...] or value}

    storage_path = get_storage_path()
    csv_files = find_csv_files(storage_path)

    # Filter based on pattern
    if pattern:
        try:
            # Limit pattern length to prevent ReDoS attacks
            if len(pattern) > 200:
                return jsonify({"error": "Pattern too long"}), 400

            # Compile pattern first to catch errors early
            compiled_pattern = re.compile(pattern, re.IGNORECASE)
            csv_files = [f for f in csv_files if compiled_pattern.search(f['path'])]
        except re.error:
            return jsonify({"error": "Invalid regex pattern"}), 400
        except Exception:
            return jsonify({"error": "Error processing pattern"}), 400

    # Filter based on parameter values
    if param_filters:
        filtered = []
        for csv_file in csv_files:
            params = csv_file.get('parameters', {})
            matches = True
            for param_name, param_values in param_filters.items():
                # Support both single value (string) and multiple values (list)
                if isinstance(param_values, list):
                    # Check if parameter value matches any of the selected values
                    if param_name not in params or str(params[param_name]) not in [str(v) for v in param_values]:
                        matches = False
                        break
                else:
                    # Legacy support for single value
                    if param_name not in params or str(params[param_name]) != str(param_values):
                        matches = False
                        break
            if matches:
                filtered.append(csv_file)
        csv_files = filtered

    return jsonify({"instances": csv_files, "count": len(csv_files)})


@app.route('/api/parameters')
def get_parameters():
    """Get aggregated parameters from all commit directories."""
    storage_path = get_storage_path()
    params = aggregate_parameters(storage_path)
    return jsonify({"parameters": params})


def run_app(host='127.0.0.1', port=5000, debug=False):
    """Run the Flask application.

    Args:
        host: Host address to bind to
        port: Port number to listen on
        debug: Enable debug mode (WARNING: should be False in production)
    """
    app.run(host=host, port=port, debug=debug)


if __name__ == '__main__':
    run_app()
