"""Stage-1 vision projector training for CORTEX AI.

Trains *only* ``VisionProjector`` on top of a frozen real vision encoder and a frozen small
causal LM::

    IMAGE -> real vision encoder -> VisionProjector -> PrefixFusion -> frozen causal LM -> logits

Nothing here touches the DeepSeek-V4-Pro shards: the ``model-*-of-*.safetensors`` files are only
read to hash their LFS pointers, and the checkpoint written by this script is a separate file.

Examples::

    python examples/train_vision_projector.py --config configs/vision_projector.json
    python examples/train_vision_projector.py --max-steps 20 --batch-size 2
"""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import asdict, fields
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from cortex_ai.vision.train import (  # noqa: E402
    ProjectorTrainingConfig,
    build_fusion_pipeline,
    count_parameters,
    iterate_caption_batches,
    run_training,
    save_projector,
)


def load_dataset(config: ProjectorTrainingConfig) -> tuple[list, list[str]]:
    """Load real images and captions from the configured Hub dataset."""
    from datasets import load_dataset as hf_load_dataset

    ds = hf_load_dataset(config.dataset, split=config.dataset_split, streaming=True)
    images, captions = [], []
    limit = config.limit or 0
    for example in ds:
        image = example.get("jpg") or example.get("image")
        text = example.get("txt") or example.get("caption") or example.get("text")
        if image is None or not text:
            continue
        if isinstance(text, list):
            text = text[0]
        images.append(image.convert("RGB"))
        captions.append(str(text))
        if limit and len(images) >= limit:
            break
    if not images:
        raise SystemExit(f"no usable image/caption pairs found in {config.dataset}")
    return images, captions


def parse_args(argv: list[str] | None = None) -> ProjectorTrainingConfig:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--config", help="JSON file of ProjectorTrainingConfig fields")
    defaults = ProjectorTrainingConfig()
    for f in fields(defaults):
        value = getattr(defaults, f.name)
        if isinstance(value, bool):
            p.add_argument(f"--{f.name.replace('_', '-')}", type=lambda s: s.lower() == "true", default=None)
        else:
            p.add_argument(
                f"--{f.name.replace('_', '-')}",
                type=type(value) if value is not None else str,
                default=None,
            )
    args = p.parse_args(argv)

    values = {}
    if args.config:
        values.update(json.loads(Path(args.config).read_text()))
    for f in fields(defaults):
        v = getattr(args, f.name.replace("-", "_"))
        if v is not None:
            values[f.name] = v
    return ProjectorTrainingConfig(**values)


def main(argv: list[str] | None = None) -> int:
    config = parse_args(argv)
    torch.manual_seed(config.seed)

    print("=" * 78)
    print("CORTEX AI - stage-1 vision projector training (experimental)")
    print("=" * 78)
    print(f"config        : {json.dumps(asdict(config), indent=2)}")

    from transformers import AutoImageProcessor, AutoModelForCausalLM, AutoTokenizer, SiglipVisionModel

    dtype = config.resolved_dtype()
    device = config.resolved_device()

    print(f"\n[1] loading real vision encoder : {config.vision_model}")
    image_processor = AutoImageProcessor.from_pretrained(config.vision_model)
    vision_encoder = SiglipVisionModel.from_pretrained(config.vision_model).to(dtype=dtype, device=device)

    print(f"[2] loading real frozen LM      : {config.language_model}")
    core = AutoModelForCausalLM.from_pretrained(config.language_model, dtype=dtype).to(device)
    tokenizer = AutoTokenizer.from_pretrained(config.language_model)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print("[3] assembling fusion pipeline (encoder + LM frozen, projector trainable)")
    fusion, vision_hidden, lm_hidden = build_fusion_pipeline(
        vision_encoder, core, config, image_processor=image_processor
    )
    projector_params = count_parameters(fusion.projector)
    print(f"    vision_hidden_size   = {vision_hidden}")
    print(f"    lm_hidden_size       = {lm_hidden}")
    print(f"    projector_type       = {config.projector_type}")
    print(f"    projector_parameters = {projector_params:,}")
    print(f"    dtype / device       = {config.dtype} / {config.device}")

    print(f"\n[4] loading real dataset        : {config.dataset} (split={config.dataset_split})")
    images, captions = load_dataset(config)
    print(f"    {len(images)} image/caption pairs")
    print(f"    sample caption: {captions[0][:70]!r}")

    batches = iterate_caption_batches(images, captions, tokenizer, image_processor, config)

    def on_step(step: int, loss: float) -> None:
        if step == 1 or step % config.log_every == 0:
            print(f"    step {step:4d}  loss={loss:.4f}")

    print(f"\n[5] training projector for up to {config.max_steps} steps (projector only)")
    report = run_training(
        fusion, batches, config, projector_params, vision_hidden, lm_hidden, on_step=on_step
    )

    path = save_projector(fusion.projector, config, report)
    report.checkpoint_path = str(path)

    print("\n" + "=" * 78)
    print("RESULT")
    print("=" * 78)
    print(report.to_json())
    print(f"\ncheckpoint written: {path}")
    print("\nPROVEN by this run:")
    print("  - real image -> real encoder -> projector -> PrefixFusion -> frozen LM executed")
    print("  - projector received gradients and its parameters changed")
    print("  - vision encoder and causal LM parameters are unchanged")
    print("\nNOT PROVEN:")
    print("  - that this projector works with the real DeepSeek-V4-Pro core")
    print("  - native multimodal understanding of the 1.65T checkpoint")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())