# -*- coding: utf-8 -*-
"""
This code references: https://github.com/yuz1wan/localWandb2cloudSwanlab, thanks to @yuz1wan.

------example.py------
from swanlab.converter import WandbLocalConverter

wb_local_converter = WandbLocalConverter()
wb_local_converter.run(root_wandb_dir="WANDB_LOCAL_DIR", wandb_run_dir="WANDB_LOCAL_RUN_DIR")

------command------
swanlab convert -t wandb-local --wb-dir ./wandb --wb-run-dir run-1234567890
"""

# GC frequency: collect garbage every N records
GC_INTERVAL = 10000

import argparse
import gc
import time
import glob
import json

try:
    import orjson as _orjson
    _json_loads = _orjson.loads
except ImportError:
    _json_loads = json.loads

import os
import re
import traceback
from typing import Optional, List

import yaml
import swanlab
from swanlab import echarts
from swanlab.log import swanlog as swl
from swanlab.data.porter import DataPorter
from swanlab.env import create_time
from swanlab.toolkit import ColumnInfo, ChartType

from rich.progress import Progress, BarColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn
from rich.progress import ProgressColumn
from rich.text import Text


class SizeColumn(ProgressColumn):
    """自定义列：显示文件大小（如 84.3M/7.61G）"""

    def __init__(self):
        super().__init__()

    def render(self, task: "Task") -> "Text":  # type: ignore
        completed = task.completed
        total = task.total
        if total is None:
            return Text("-/-", style="progress.filesize")
        return Text(f"{_format_size(int(completed))}/{_format_size(int(total))}", style="progress.filesize")


def _format_size(size_bytes: int) -> str:
    """格式化字节为人类可读格式"""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size_bytes < 1024:
            return f"{size_bytes:.1f}{unit}"
        size_bytes /= 1024
    return f"{size_bytes:.1f}PB"

try:
    from wandb.sdk.internal.datastore import DataStore
    from wandb.proto import wandb_internal_pb2
except ImportError:
    # We add a print here because if this fails, swanlog might not be available.
    print("Error: Wandb is required to parse local log files. Please install it with 'pip install wandb'.")
    exit(1)


class WandbLocalConverter:
    """
    A robust converter to migrate local wandb run data to SwanLab.

    This class reliably parses all data from a .wandb file by using the
    `ds.scan_data()` method and converting Protobuf messages to dictionaries
    with `MessageToDict`.
    """
    def __init__(
        self,
        project: Optional[str] = None,
        workspace: Optional[str] = None,
        mode: Optional[str] = "cloud",
        tags: Optional[List[str]] = None,
        logdir: Optional[str] = None,
        **kwargs,
    ):
        """
        Initializes the converter with SwanLab project details.

        Args:
            project (Optional[str]): The SwanLab project name. If None, it will be inferred from the wandb run.
            workspace (Optional[str]): The SwanLab workspace.
            logdir (Optional[str]): The local directory to save SwanLab logs.
        """
        self.project = project
        self.workspace = workspace
        self.mode = mode
        self.tags = tags
        self.logdir = logdir

    def _find_run_dirs(self, root_wandb_dir: str, wandb_run_dir: Optional[str] = None) -> List[str]:
        """Finds all wandb run directories within a given root directory."""
        if wandb_run_dir:
            patterns = [
                os.path.join(root_wandb_dir, wandb_run_dir)
            ]
            found_path = os.path.join(root_wandb_dir, wandb_run_dir)
        else:
            patterns = [
                os.path.join(root_wandb_dir, "run-*/"),
                os.path.join(root_wandb_dir, "offline-run-*/")
            ]
            found_path = root_wandb_dir

        run_dirs = []
        for pattern in patterns:
            run_dirs.extend(glob.glob(pattern))
        if not run_dirs:
            swl.warning(f"No wandb run directories found in '{found_path}'.")
        return run_dirs

    def _proto_items_to_dict(self, items) -> dict:
        """Directly read proto repeated item fields (key/nested_key/value_json) into a dict.

        This avoids the expensive MessageToDict conversion by accessing proto fields directly.
        """
        mapping = {}
        for item in items:
            key = item.key or '/'.join(item.nested_key)
            if not key:
                continue
            try:
                mapping[key] = _json_loads(item.value_json)
            except (ValueError, Exception):
                swl.warning(f"Could not decode json for key '{key}': {item.value_json}")
        return mapping

    def _validate_path(self, base_dir: str, file_path: str) -> Optional[str]:
        """Validate that file_path stays within base_dir to prevent path traversal."""
        if not file_path:
            return None
        abs_base = os.path.abspath(base_dir)
        abs_path = os.path.abspath(os.path.join(base_dir, file_path))
        if not abs_path.startswith(abs_base + os.sep) and abs_path != abs_base:
            swl.warning(f"Path traversal attempt detected: {file_path}")
            return None
        return abs_path

    def _filter_text_columns(self, columns, data):
        """Filter out non-text columns (images, media) from table data."""
        non_text_indices = {
            i for row in data for i, cell in enumerate(row) if isinstance(cell, dict) and "_type" in cell
        }
        text_indices = [i for i in range(len(columns)) if i not in non_text_indices]
        return (
            [columns[i] for i in text_indices],
            [[row[i] for i in text_indices if i < len(row)] for row in data]
        )

    def _parse_run(self, run_dir: str):
        """Parses a single wandb run directory and converts it to a SwanLab run."""
        wandb_files = glob.glob(os.path.join(run_dir, "*.wandb"))
        if not wandb_files:
            swl.error(f"No .wandb file found in {run_dir}. Skipping.")
            return

        wandb_file_path = wandb_files[0]
        files_root_dir = os.path.join(run_dir, "files")
        run_basename = os.path.basename(run_dir.rstrip('/\\'))

        # 获取文件大小用于进度条
        wandb_file_size = os.path.getsize(wandb_file_path)

        ds = DataStore()
        ds.open_for_scan(wandb_file_path)

        swanlab_run = None
        run_metadata = {'id': None, 'name': None, 'notes': None, 'project': None}
        run_config = {}

        # Direct-write state (bypasses swanlab_run.log() overhead for scalar metrics)
        porter = None           # set after swanlab.init()
        column_kids = {}  # key -> kid str
        epoch_counters = {}  # key -> cumulative epoch count
        upload_pre_count = 0  # items uploaded before upload progress bar appears
        total_scalars_logged = 0  # total scalar data points logged

        def _log_scalars_direct(scalars, step):
            """Write float scalars directly to the porter, bypassing log() overhead."""
            nonlocal total_scalars_logged
            if not scalars or porter is None:
                return
            for key in scalars:
                if key not in column_kids:
                    kid = len(column_kids)
                    column_kids[key] = str(kid)
                    split_key = key.split("/")
                    sname = split_key[0] if len(split_key) > 1 and split_key[0] else None
                    porter.trace_column(ColumnInfo(
                        key=key, kid=str(kid), name=key, cls='CUSTOM',
                        chart_type=ChartType.LINE, chart_reference='STEP',
                        section_name=sname, section_type="PUBLIC",
                    ))
                epoch_counters[key] = epoch_counters.get(key, 0) + 1
            total_scalars_logged += len(scalars)
            porter.trace_scalars_step(step, scalars, dict(epoch_counters), create_time())

        def _finish_with_progress():
            """Run swanlab_run.finish() while showing a Rich upload progress bar."""
            _pool = porter._pool if porter is not None else None
            if _pool is None:
                swanlab_run.finish()
                return
            total = len(column_kids) + total_scalars_logged
            up = Progress(
                TextColumn("[bold green]{task.description}"),
                BarColumn(bar_width=40),
                TextColumn("{task.completed}/{task.total} items"),
                TimeElapsedColumn(),
                TimeRemainingColumn(),
            )
            up.start()
            t = up.add_task("Uploading to SwanLab", total=total, completed=upload_pre_count)

            last_completed = upload_pre_count
            stall_check_time = time.time()

            def _upload_cb(n):
                nonlocal last_completed, stall_check_time
                last_completed += n
                stall_check_time = time.time()
                up.update(t, advance=n)

            _pool.collector.upload_callback = _upload_cb

            # Monitor for stalls and update description
            import threading
            stop_monitor = False
            def _monitor_stalls():
                nonlocal stop_monitor
                while not stop_monitor:
                    time.sleep(1)
                    if stop_monitor:
                        break
                    elapsed = time.time() - stall_check_time
                    if elapsed > 5 and last_completed < total:
                        remaining = total - last_completed
                        batch_info = f"batch ~{min(1000, remaining)} items" if remaining > 0 else "final batch"
                        up.update(t, description=f"[bold yellow]Uploading to SwanLab (processing {batch_info}, {int(elapsed)}s)")
                    elif last_completed < total:
                        up.update(t, description="[bold green]Uploading to SwanLab")

            monitor_thread = threading.Thread(target=_monitor_stalls, daemon=True)
            monitor_thread.start()

            # Suppress swanlog.info during finish() to hide the "View project/run" URL reprints
            _orig_info = swl.info
            swl.info = lambda *a, **k: None
            try:
                swanlab_run.finish()
            finally:
                swl.info = _orig_info
                stop_monitor = True
            up.update(t, completed=total, description="[bold green]Uploading to SwanLab")
            up.stop()

        def initialize_swanlab_run_if_needed():
            """Lazy initializer for the SwanLab run."""
            nonlocal swanlab_run, progress, task_id, porter
            if swanlab_run is not None:
                return

            # 暂停进度条，避免与 swanlab.init 中的 rich Status 冲突
            current_progress = 0
            if progress is not None and task_id is not None:
                current_progress = int(progress.tasks[task_id].completed)
                progress.stop()

            if run_metadata['id'] is None:
                match = re.search(r"-([a-zA-Z0-9]+)$", os.path.basename(run_dir.rstrip('/\\')))
                if match:
                    run_metadata['id'] = match.group(1)
                else:
                    raise RuntimeError(f"Fatal: Cannot determine run_id for {run_dir}.")
            if run_metadata['name'] is None:
                run_metadata['name'] = run_metadata['id']

            swl.info(f"Converting run: {run_metadata['name']} (ID: {run_metadata['id']})")
            swanlab_run = swanlab.init(
                project=self.project or run_metadata['project'] or "wandb-imports",
                workspace=self.workspace,
                experiment_name=run_metadata['name'],
                description=run_metadata['notes'],
                logdir=self.logdir,
                mode=self.mode,
                tags=self.tags,
            )
            porter = DataPorter._instance
            # Set pre-count callback: track items uploaded during parse before the upload bar appears
            if porter is not None and porter._pool is not None:
                def _pre_upload_cb(n):
                    nonlocal upload_pre_count
                    upload_pre_count += n
                porter._pool.collector.upload_callback = _pre_upload_cb

            # 恢复进度条
            if progress is not None and task_id is not None:
                try:
                    progress.start()
                    progress.update(task_id, completed=current_progress)
                except Exception:
                    # 如果 start 失败，重新创建进度条并恢复进度
                    progress = Progress(
                        TextColumn("[bold blue]{task.description}"),
                        BarColumn(bar_width=40),
                        SizeColumn(),
                        TimeElapsedColumn(),
                        TimeRemainingColumn(),
                    )
                    progress.start()
                    task_id = progress.add_task(
                        f"Parsing  {run_basename}",
                        total=wandb_file_size,
                        completed=current_progress
                    )

            try:
                config_path = os.path.join(files_root_dir, "config.yaml")
                with open(config_path, 'r', encoding='utf-8') as f:
                    wandb_config_yaml = yaml.safe_load(f)
                    cleaned_config = {k: v['value'] for k, v in wandb_config_yaml.items() if isinstance(v, dict) and 'value' in v}
                    run_config.update(cleaned_config)
            except FileNotFoundError:
                pass
            
            final_config = {"wandb_run_id": run_metadata['id'], "wandb_local_path": run_dir}
            final_config.update(run_config)
            swanlab_run.config.update(final_config)

        # Core Logic: Scan the .wandb file record by record
        # 初始化进度条
        progress = None
        task_id = None
        record_count = 0

        if Progress is not None:
            progress = Progress(
                TextColumn("[bold blue]{task.description}"),
                BarColumn(bar_width=40),
                SizeColumn(),
                TimeElapsedColumn(),
                TimeRemainingColumn(),
            )
            progress.start()
            task_id = progress.add_task(f"Parsing  {run_basename}", total=wandb_file_size)

        while True:
            record_bin = ds.scan_data()
            if record_bin is None:
                break

            # 更新进度条
            if progress is not None and task_id is not None:
                progress.update(task_id, advance=len(record_bin))

            record_pb = wandb_internal_pb2.Record()
            record_pb.ParseFromString(record_bin)

            record_type = record_pb.WhichOneof("record_type")

            # Process record by directly accessing protobuf fields (no MessageToDict)
            if record_type == "run":
                run = record_pb.run
                run_metadata.update({
                    'id': run.run_id or None,
                    'name': run.display_name or None,
                    'notes': run.notes or None,
                    'project': run.project or None,
                })
                run_config.update(self._proto_items_to_dict(run.config.update))
            elif record_type == "config":
                run_config.update(self._proto_items_to_dict(record_pb.config.update))
            elif record_type == "history":
                initialize_swanlab_run_if_needed()
                scalar_dict = {}
                media_dict = {}
                grouped_items = {}
                step = 0
                # First pass: group items by base key
                for item in record_pb.history.item:
                    key = item.key or '/'.join(item.nested_key)
                    if not key:
                        continue
                    value_json = item.value_json
                    if key == '_step':
                        try:
                            step = int(float(value_json))
                        except (ValueError, TypeError):
                            pass
                        continue
                    if key.startswith('_'):
                        continue
                    # Check if key has nested structure (e.g., "table/_type")
                    if '/' in key:
                        base_key, sub_key = key.split('/', 1)
                        if base_key not in grouped_items:
                            grouped_items[base_key] = {}
                        try:
                            grouped_items[base_key][sub_key] = _json_loads(value_json)
                        except (ValueError, Exception):
                            grouped_items[base_key][sub_key] = value_json
                    else:
                        # Fast path: direct float conversion for scalars
                        try:
                            scalar_dict[key] = float(value_json)
                            continue
                        except (ValueError, TypeError):
                            pass
                        # Slow path: full JSON parse
                        try:
                            value = _json_loads(value_json)
                            if isinstance(value, int):
                                scalar_dict[key] = float(value)
                            elif isinstance(value, dict) and "path" in value:
                                validated_path = self._validate_path(files_root_dir, value["path"])
                                if validated_path and os.path.exists(validated_path):
                                    if value.get("_type") == "image-file":
                                        media_dict[key] = swanlab.Image(validated_path)
                                    elif value.get("_type") == "audio-file":
                                        media_dict[key] = swanlab.Audio(validated_path)
                        except (ValueError, Exception):
                            pass

                # Second pass: process grouped items for tables and media
                for base_key, props in grouped_items.items():
                    if props.get('_type') == 'table-file' and 'path' in props:
                        validated_path = self._validate_path(files_root_dir, props['path'])
                        if validated_path and os.path.exists(validated_path):
                            try:
                                with open(validated_path, 'r', encoding='utf-8') as f:
                                    table_data = _json_loads(f.read())
                                columns = table_data.get("columns", [])
                                data = table_data.get("data", [])
                                # Filter text-only columns
                                filtered_cols, filtered_data = self._filter_text_columns(columns, data)
                                if filtered_cols:
                                    table = echarts.Table()
                                    table.add(filtered_cols, filtered_data)
                                    media_dict[base_key] = table
                            except Exception as e:
                                swl.warning(f"Failed to parse table from {validated_path}: {e}")
                    elif props.get('_type') == 'image-file' and 'path' in props:
                        validated_path = self._validate_path(files_root_dir, props['path'])
                        if validated_path and os.path.exists(validated_path):
                            media_dict[base_key] = swanlab.Image(validated_path)
                    elif props.get('_type') == 'audio-file' and 'path' in props:
                        validated_path = self._validate_path(files_root_dir, props['path'])
                        if validated_path and os.path.exists(validated_path):
                            media_dict[base_key] = swanlab.Audio(validated_path)

                if scalar_dict or media_dict:
                    if scalar_dict:
                        _log_scalars_direct(scalar_dict, step)
                    if media_dict:
                        swanlab_run.log(media_dict, step=step)
                del scalar_dict, media_dict

            # GC every GC_INTERVAL records to reduce overhead
            record_count += 1
            if record_count % GC_INTERVAL == 0:
                gc.collect()

        # 关闭进度条
        if progress is not None:
            progress.stop()

        if swanlab_run:
            swl.info(f"Finished Parsing run: {run_metadata['name']}")
            _finish_with_progress()
        else:
            try:
                initialize_swanlab_run_if_needed()
                if swanlab_run:
                    swl.warning(f"Run in {run_dir} has no metrics, but its config was saved.")
                    _finish_with_progress()
            except RuntimeError as e:
                swl.error(str(e))

        # 清理 DataStore 对象，释放文件句柄
        del ds
        gc.collect()

    def run(self, root_wandb_dir: str, wandb_run_dir: Optional[str] = None):
        """The main entry point to start the conversion process."""
        swl.info(f"Starting import from wandb directory: {os.path.abspath(root_wandb_dir)}")
        run_dirs = self._find_run_dirs(root_wandb_dir, wandb_run_dir)
        
        if not run_dirs:
            swl.info("No runs found to import. Exiting.")
            return

        swl.info(f"Found {len(run_dirs)} runs to import.")
        
        for i, run_dir in enumerate(run_dirs):
            print("\n")
            swl.info(f"--- Processing run {i+1}/{len(run_dirs)} ---")
            try:
                self._parse_run(run_dir)
            except Exception as e:
                swl.error(f"FATAL: Failed to convert run in {run_dir}: {e}")
                traceback.print_exc()
        
        swl.info("="*20)
        swl.info("All runs processed.")
        swl.info("="*20)


def main():
    """Main function to parse arguments and run the converter."""
    parser = argparse.ArgumentParser(
        description="Import local Wandb runs into SwanLab.",
        formatter_class=argparse.RawTextHelpFormatter
    )
    
    parser.add_argument(
        "--wandb_dir",
        "-d",
        type=str,
        default="./wandb",
        required=True,
        help="Path to the local 'wandb' directory containing run folders (e.g., './wandb')."
    )
    
    parser.add_argument(
        "--wandb_run_dir",
        "-r",
        type=str,
        default=None,
        help="Path to the local 'wandb' directory containing run folders (e.g., './wandb')."
    )
    
    parser.add_argument(
        "--project",
        "-p",
        type=str,
        default=None,
        help="Specify the SwanLab project name. If not set, it's inferred from the wandb run."
    )
    
    parser.add_argument(
        "--workspace",
        "-w",
        type=str,
        default=None,
        help="Specify the SwanLab workspace."
    )
    
    parser.add_argument(
        "--logdir",
        "-l",
        type=str,
        default=None,
        help="Specify the root directory for SwanLab logging. Defaults to './swanlog'."
    )
    
    args = parser.parse_args()
    
    swl.info("SwanLab Importer for Wandb")
    swl.info("---------------------------")
    swl.info(f"Source Wandb Directory: {os.path.abspath(args.wandb_dir)}")
    swl.info(f"Source Wandb Run Directory: {args.wandb_run_dir or '(all runs)'}")
    swl.info(f"Target SwanLab Project: {args.project or '(inferred from run)'}")
    swl.info(f"Target SwanLab Workspace: {args.workspace or '(default)'}")
    swl.info(f"SwanLab Log Directory: {args.logdir or './swanlog'}")
    swl.info("---------------------------")

    converter = WandbLocalConverter(
        project=args.project,
        workspace=args.workspace,
        logdir=args.logdir
    )
    converter.run(root_wandb_dir=args.wandb_dir, wandb_run_dir=args.wandb_run_dir)


if __name__ == '__main__':
    main()