# scripts/train_gemma3.py

import os
import sys
import json
import torch
from pathlib import Path
from datasets import load_dataset
from unsloth import FastModel
from trl import SFTTrainer, SFTConfig
from unsloth.chat_templates import get_chat_template, train_on_responses_only

# PyTorchメモリ最適化
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

def setup_environment():
    """環境設定"""
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    # float16の自動混合精度を有効化
    torch.set_float32_matmul_precision('medium')

    # バージョン確認
    try:
        import unsloth
        print(f"Unsloth version: {unsloth.__version__}")
    except ImportError:
        print("Unslothバージョン取得失敗")

    import transformers
    print(f"Transformers version: {transformers.__version__}")

    # GPU精度サポート確認
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        if "RTX" in gpu_name or "Tesla" in gpu_name or "A100" in gpu_name:
            print(f"GPU {gpu_name} supports bfloat16")
        else:
            print(f"GPU {gpu_name} precision support needs verification")

def get_gpu_optimal_config():
    """GPU最適設定を決定"""
    if not torch.cuda.is_available():
        raise RuntimeError("CUDAが利用できません。GPUが必要です。")

    gpu_memory_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
    device_name = torch.cuda.get_device_name(0)

    print(f"GPU: {device_name} ({gpu_memory_gb:.1f}GB)")

    config = {
        "model_name": "unsloth/gemma-3-4b-it",
        "model_size": "4B",
        "max_seq_length": 1024,
        "batch_size": 1,
        "gradient_accumulation": 4,
        "lora_rank": 8,
        "lora_alpha": 8
    }

    return config

def load_model_and_tokenizer(config):
    """Unsloth公式準拠のメモリ効率的モデル読み込み"""
    print(f"モデル読み込み中: {config['model_name']}")

    model, tokenizer = FastModel.from_pretrained(
        model_name=config["model_name"],
        max_seq_length=config["max_seq_length"],
        load_in_4bit=True,
        dtype=None,
        trust_remote_code=True,
        device_map={"": 0},
        attn_implementation="eager"
    )

    tokenizer = get_chat_template(
        tokenizer,
        chat_template="gemma-3",
    )

    print("モデル読み込み完了")
    return model, tokenizer

def setup_lora(model, config):
    """Unsloth公式準拠LoRA設定"""
    print("LoRA設定中...")

    model = FastModel.get_peft_model(
        model,
        r=config["lora_rank"],
        lora_alpha=config["lora_alpha"],
        lora_dropout=0,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=3407,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
    )

    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())

    print(f"学習可能パラメータ: {trainable_params:,} / {total_params:,} ({trainable_params/total_params*100:.2f}%)")

    return model

def prepare_dataset(tokenizer, data_path="training_data/code_review_data.jsonl"):
    """
    Unsloth公式Gemma3準拠のデータセット準備
    instruction/input/output → conversations → text (chat template applied)
    """
    print(f"データセット読み込み: {data_path}")

    # データファイル存在確認
    if not Path(data_path).exists():
        raise FileNotFoundError(f"データファイルが見つかりません: {data_path}")

    # JSONL形式のデータセット読み込み
    dataset = load_dataset("json", data_files=data_path, split="train")
    print(f"データ件数: {len(dataset)}")

    # データ形式確認
    if len(dataset) > 0:
        sample = dataset[0]
        required_fields = ["instruction", "input", "output"]
        for field in required_fields:
            if field not in sample:
                raise ValueError(f"必須フィールドが見つかりません: {field}")

        print("データ形式確認完了")

    def convert_to_conversations(examples):
        """
        Unsloth公式準拠: instruction/input/output → conversations形式
        """
        instructions = examples["instruction"]
        inputs = examples["input"]
        outputs = examples["output"]

        conversations_list = []

        # バッチ処理
        if isinstance(instructions, list):
            for inst, inp, out in zip(instructions, inputs, outputs):
                # ユーザープロンプト作成
                if inp and inp.strip():
                    user_content = f"{inst}\n\n{inp}"
                else:
                    user_content = inst

                # Unsloth conversations形式
                conversations = [
                    {"role": "user", "content": user_content},
                    {"role": "assistant", "content": out}
                ]
                conversations_list.append(conversations)
        else:
            # 単一処理
            if inputs and inputs.strip():
                user_content = f"{instructions}\n\n{inputs}"
            else:
                user_content = instructions

            conversations = [
                {"role": "user", "content": user_content},
                {"role": "assistant", "content": outputs}
            ]
            conversations_list = [conversations]

        return {"conversations": conversations_list}

    # conversations形式に変換
    print("Unsloth conversations形式に変換中...")
    dataset = dataset.map(convert_to_conversations, batched=True)

    # Unsloth公式準拠: チャットテンプレート適用
    def formatting_prompts_func(examples):
        """
        Unsloth公式Gemma3準拠のフォーマット関数
        conversations → text (Gemma3 chat template + <bos>除去)
        """
        convos = examples["conversations"]
        texts = [
            tokenizer.apply_chat_template(
                convo,
                tokenize=False,
                add_generation_prompt=False
            ).removeprefix('<bos>')
            for convo in convos
        ]
        return {"text": texts}

    print("Gemma3チャットテンプレート適用中...")
    dataset = dataset.map(formatting_prompts_func, batched=True)

    # 元のカラムを削除（textのみ残す）
    columns_to_remove = [col for col in dataset.column_names if col != "text"]
    if columns_to_remove:
        dataset = dataset.remove_columns(columns_to_remove)

    # 最終的なデータセット構造確認
    print(f"最終データ件数: {len(dataset)}")

    # 変換結果確認
    if len(dataset) > 0:
        sample = dataset[0]
        if not sample['text'].startswith('<start_of_turn>'):
            print(f"想定外の開始トークン: {sample['text'][:50]}")

        if not isinstance(sample['text'], str):
            raise ValueError(f"データ型エラー: text={type(sample['text'])}")

    print("データセット準備完了")
    return dataset

def create_training_config(config, output_dir="fine_tuned_gemma3"):
    """Unsloth最適化学習設定"""
    training_config = SFTConfig(
        dataset_text_field="text",
        per_device_train_batch_size=config["batch_size"],
        gradient_accumulation_steps=config["gradient_accumulation"],
        warmup_steps=5,
        max_steps=20,
        learning_rate=2e-4,

        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",

        fp16=False,
        bf16=True,
        dataloader_pin_memory=False,

        logging_steps=1,
        output_dir=output_dir,
        save_steps=20,
        save_strategy="steps",
        save_total_limit=1,

        seed=3407,
        report_to="none",
        dataset_num_proc=2
    )

    training_config.max_seq_length = config["max_seq_length"]

    return training_config

def train_model(model, tokenizer, dataset, training_args):
    """Unsloth学習実行"""
    print("学習開始...")

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset,
        eval_dataset=None,
        args=training_args,
        dataset_text_field="text",
        max_seq_length=getattr(training_args, 'max_seq_length', 1024),
        packing=False,
    )

    # train_on_responses_onlyを適用
    trainer = train_on_responses_only(
        trainer,
        instruction_part="<start_of_turn>user\n",
        response_part="<start_of_turn>model\n",
    )

    torch.cuda.empty_cache()

    try:
        trainer_stats = trainer.train()

        print("学習完了")
        print(f"最終損失: {trainer_stats.training_loss:.4f}")
        print(f"学習ステップ数: {trainer_stats.global_step}")

    except RuntimeError as e:
        if "out of memory" in str(e).lower():
            print("GPU メモリ不足: バッチサイズやLoRAランクを削減してください")
        raise
    finally:
        torch.cuda.empty_cache()

    return trainer

# create_chat_prompt関数は削除（Unsloth公式準拠の前処理に統合済み）

def save_model(model, tokenizer, config, output_dir):
    """
    CI/CD対応モデル保存
    推論スクリプトが正確に読み込めるように設定情報も保存
    """
    print(f"モデル保存中: {output_dir}")

    # 出力ディレクトリ作成
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    # LoRAアダプター保存
    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)

    # 推論用設定情報保存
    model_config = {
        "model_type": "unsloth_lora",
        "base_model_name": config["model_name"],
        "base_model_size": config["model_size"],
        "max_seq_length": config["max_seq_length"],
        "lora_rank": config["lora_rank"],
        "lora_alpha": config["lora_alpha"],
        "chat_template": "gemma-3",
        "training_completed": True,
        "inference_ready": True
    }

    with open(f"{output_dir}/model_config.json", 'w') as f:
        json.dump(model_config, f, indent=2)

    print(f"LoRAアダプター保存完了: {output_dir}")

    # 推論用マージ済みモデル保存（オプション）
    merged_dir = f"{output_dir}_merged"
    try:
        model.save_pretrained_merged(merged_dir, tokenizer)

        # マージ済みモデル用設定
        merged_config = model_config.copy()
        merged_config["model_type"] = "merged_model"

        with open(f"{merged_dir}/model_config.json", 'w') as f:
            json.dump(merged_config, f, indent=2)

        print(f"マージ済みモデル保存完了: {merged_dir}")
    except Exception as e:
        print(f"マージ済みモデル保存エラー: {e}")

    # 保存結果サマリー
    print(f"CI/CDアップロード対象: {output_dir}/")
    if Path(merged_dir).exists():
        print(f"推論用モデル: {merged_dir}/")

def main():
    """メイン処理"""
    print("=== Gemma 3 + Unsloth ファインチューニング ===")

    try:
        # 環境設定
        setup_environment()

        # GPU最適設定取得
        config = get_gpu_optimal_config()

        # モデル読み込み
        model, tokenizer = load_model_and_tokenizer(config)

        # LoRA設定
        model = setup_lora(model, config)

        # データセット準備
        dataset = prepare_dataset(tokenizer)

        # 学習設定
        training_config = create_training_config(config)

        # 学習実行
        trainer = train_model(model, tokenizer, dataset, training_config)

        # モデル保存
        save_model(model, tokenizer, config, training_config.output_dir)

        # メモリ使用量表示
        if torch.cuda.is_available():
            memory_used = torch.cuda.max_memory_allocated() / 1e9
            print(f"最大GPU使用量: {memory_used:.2f}GB")

        print("すべての処理が完了しました")

    except Exception as e:
        print(f"エラー: {e}")
        print(f"エラータイプ: {type(e).__name__}")
        import traceback
        print("詳細なスタックトレース:")
        traceback.print_exc()
        raise
    finally:
        # GPUメモリクリア
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

if __name__ == "__main__":
    main()

# scripts/code_reviewer.py
from unsloth import FastLanguageModel
import torch
import sys
import os

def load_fine_tuned_model():
    """ファインチューニング済みモデルを読み込み"""
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name="fine_tuned_gemma3",
        max_seq_length=2048,
        dtype=None,
        load_in_4bit=True,
    )
    # 推論モードに設定（メモリ効率化・高速化）
    FastLanguageModel.for_inference(model)
    return model, tokenizer

def review_code(code_diff, model, tokenizer):
    """コードの差分をレビューする"""
    prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
以下のコード変更をレビューしてください。改善点があれば具体的に指摘してください。

### Input:
{code_diff}

### Response:"""

    inputs = tokenizer([prompt], return_tensors="pt").to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.15,
        )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Response:以降を抽出
    if "### Response:" in response:
        return response.split("### Response:")[-1].strip()
    return response

def review_from_gitlab_diff():
    """GitLab CI環境でgit diffからレビューを実行"""
    print("AI Code Review")

    # git diffの取得
    import subprocess
    try:
        # 最新のコミットの差分を取得
        diff_output = subprocess.check_output(
            ["git", "diff", "HEAD~1", "HEAD"],
            text=True
        )

        if not diff_output.strip():
            print("変更が検出されませんでした")
            return

    except subprocess.CalledProcessError:
        print("エラー: git diffの取得に失敗しました")
        return

    # モデルの読み込み
    model, tokenizer = load_fine_tuned_model()

    # コードレビューの実行
    print("レビューを生成中")
    review = review_code(diff_output, model, tokenizer)

    print("【レビュー結果】")
    print("-" * 50)
    print(review)
    print("-" * 50)

    # 結果をファイルに保存（CI/CDアーティファクト用）
    with open("review_result.txt", "w", encoding="utf-8") as f:
        f.write(f"AI Code Review Result\n")
        f.write(f"{'=' * 50}\n\n")
        f.write(review)

def main():
    if len(sys.argv) > 1:
        # ファイルパスが指定された場合
        diff_file = sys.argv[1]
        if not os.path.exists(diff_file):
            print(f"エラー: ファイルが見つかりません: {diff_file}")
            sys.exit(1)

        with open(diff_file, 'r', encoding='utf-8') as f:
            code_diff = f.read()

        model, tokenizer = load_fine_tuned_model()
        review = review_code(code_diff, model, tokenizer)

        print("AI Code Review")
        print(review)
    else:
        # 引数なしの場合はGitLab CI環境を想定
        review_from_gitlab_diff()

if __name__ == "__main__":
    main()