import random
import traceback
from pathlib import Path

import torch
import torch.nn as nn
from PIL import Image
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
from transformers import (
    LayoutLMv3Processor,
    LayoutLMv3ForSequenceClassification,
    get_scheduler,
)
from tqdm import tqdm
from loguru import logger

import config
from train.schemas import TrainRequest, TrainStatus
from utils.cache import (
    load_cache, save_label2id,
    load_replay_buffer, update_replay_buffer, load_label2id
)

# 학습 상태 전역
train_status = TrainStatus(
    running=False, epoch=0, total=0, loss=0.0, message="대기 중"
)
stop_requested = False


# ── Dataset ───────────────────────────────────────
class InvoiceDataset(Dataset):
    def __init__(self, samples: list, processor):
        self.samples   = samples
        self.processor = processor

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        s        = self.samples[idx]
        image    = Image.open(s["image_path"]).convert("RGB")
        encoding = self.processor(
            image, text=s["words"], boxes=s["boxes"],
            return_tensors="pt", truncation=True,
            padding="max_length", max_length=config.MAX_LEN,
        )
        return {
            "input_ids":      encoding["input_ids"].squeeze(),
            "attention_mask": encoding["attention_mask"].squeeze(),
            "bbox":           encoding["bbox"].squeeze(),
            "pixel_values":   encoding["pixel_values"].squeeze(),
            "labels":         torch.tensor(s["label"], dtype=torch.long),
        }


# ── Fisher (EWC) ──────────────────────────────────
def compute_fisher(model, loader) -> dict:
    fisher = {n: torch.zeros_like(p)
              for n, p in model.named_parameters() if p.requires_grad}
    model.eval()
    for batch in loader:
        model.zero_grad()
        out = model(
            input_ids      = batch["input_ids"].to(config.DEVICE),
            attention_mask = batch["attention_mask"].to(config.DEVICE),
            bbox           = batch["bbox"].to(config.DEVICE),
            pixel_values   = batch["pixel_values"].to(config.DEVICE),
            labels         = batch["labels"].to(config.DEVICE),
        )
        out.loss.backward()
        for n, p in model.named_parameters():
            if p.requires_grad and p.grad is not None:
                fisher[n] += p.grad.pow(2)
    return {n: f / max(len(loader), 1) for n, f in fisher.items()}


# ── 검증 ──────────────────────────────────────────
def evaluate_loader(model, loader) -> float:
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for batch in loader:
            out = model(
                input_ids      = batch["input_ids"].to(config.DEVICE),
                attention_mask = batch["attention_mask"].to(config.DEVICE),
                bbox           = batch["bbox"].to(config.DEVICE),
                pixel_values   = batch["pixel_values"].to(config.DEVICE),
            )
            preds   = out.logits.argmax(dim=-1).cpu()
            correct += (preds == batch["labels"]).sum().item()
            total   += len(batch["labels"])
    return correct / total if total else 0.0


def _get_latest_model_path():
    """
    날짜/차수 폴더에서 가장 최신 모델 경로 반환
    files/models/20260305/0002/ ← 최신
    """
    base = config.MODEL_BASE_PATH
    all_models = sorted([
        d for d in base.rglob("config.json")
    ])
    if not all_models:
        return None

    print(all_models[-1].parent)
    return all_models[-1].parent


# ── 학습 루프 ─────────────────────────────────────
def run_train_loop(model, processor, train_s, val_s, req: TrainRequest,
                   label2id: dict, strategy: str, save_path: Path):
    global train_status

    train_loader = DataLoader(
        InvoiceDataset(train_s, processor),
        batch_size=req.batch_size, shuffle=True
    )
    val_loader = DataLoader(
        InvoiceDataset(val_s, processor),
        batch_size=req.batch_size
    )
    optimizer = AdamW(model.parameters(), lr=req.lr)
    scheduler = get_scheduler(
        "cosine", optimizer=optimizer,
        num_warmup_steps=len(train_loader),
        num_training_steps=len(train_loader) * req.epochs,
    )

    # EWC 사전 계산
    ewc_params, ewc_fisher = {}, {}
    if strategy == "ewc":
        ewc_params = {n: p.clone().detach()
                      for n, p in model.named_parameters() if p.requires_grad}
        ewc_fisher = compute_fisher(model, val_loader)
        logger.info("EWC Fisher 계산 완료")

    best_acc      = 0.0
    total_batches = len(train_loader) * req.epochs

    for epoch in range(req.epochs):
        model.train()
        total_loss = 0

        pbar = tqdm(
            train_loader,
            desc=f"Epoch {epoch+1}/{req.epochs}",
            ncols=100,
            leave=True,
        )
        for batch_idx, batch in enumerate(pbar):
            if stop_requested:
                logger.info("학습 중단 요청 감지")
                return

            out = model(
                input_ids      = batch["input_ids"].to(config.DEVICE),
                attention_mask = batch["attention_mask"].to(config.DEVICE),
                bbox           = batch["bbox"].to(config.DEVICE),
                pixel_values   = batch["pixel_values"].to(config.DEVICE),
                labels         = batch["labels"].to(config.DEVICE),
            )
            loss = out.loss
            if strategy == "ewc" and ewc_fisher:
                for n, p in model.named_parameters():
                    if n in ewc_fisher:
                        loss += req.ewc_lambda * (
                            ewc_fisher[n] * (p - ewc_params[n]).pow(2)
                        ).sum()

            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            total_loss += loss.item()

            pbar.set_postfix(loss=f"{loss.item():.4f}")

            # 배치 단위 상태 업데이트
            done_batches = epoch * len(train_loader) + (batch_idx + 1)
            train_status.epoch         = epoch + 1
            train_status.total         = req.epochs
            train_status.loss          = round(loss.item(), 4)
            train_status.batch         = done_batches        # 완료된 전체 배치
            train_status.total_batches = total_batches       # 전체 배치 수
            train_status.message       = (
                f"Epoch {epoch+1}/{req.epochs} | "
                f"batch {batch_idx+1}/{len(train_loader)} | "
                f"loss={loss.item():.4f}"
            )

        avg_loss = total_loss / len(train_loader)
        val_acc  = evaluate_loader(model, val_loader)

        msg = (
            f"Epoch {epoch+1}/{req.epochs} | "
            f"loss={avg_loss:.4f} | val_acc={val_acc:.4f}"
        )
        tqdm.write(f"[결과] {msg}")
        logger.info(msg)

        train_status.epoch   = epoch + 1
        train_status.loss    = round(avg_loss, 4)
        train_status.val_acc = round(val_acc, 4)
        train_status.message = msg

        if val_acc > best_acc:
            best_acc = val_acc
            model.save_pretrained(str(save_path))
            processor.save_pretrained(str(save_path))
            save_label2id(label2id)
            logger.info(f"모델 저장 (val_acc={val_acc:.4f}) → {save_path}")


# ── 학습 진입점 ───────────────────────────────────
def run_train(req: TrainRequest):
    global train_status
    stop_requested = False
    train_status = TrainStatus(
        running=True, epoch=0, total=req.epochs,
        loss=0.0, message="학습 준비 중"
    )
    try:
        from utils.cache import load_label2id

        # 날짜/차수 폴더 생성
        save_path = config.get_model_save_path()
        logger.info(f"모델 저장 경로: {save_path}")
        train_status.message = f"저장 경로: {save_path}"

        samples  = load_cache()
        label2id = load_label2id()
        id2label = {v: k for k, v in label2id.items()}

        processor_path = (
            str(config.MODEL_BASE_PATH / "preprocessor_config.json")
            if (config.MODEL_BASE_PATH / "preprocessor_config.json").exists()
            else config.MODEL_NAME
        )
        # 최신 모델 폴더 찾기 (이어 학습)
        latest_model = _get_latest_model_path()
        model_path   = str(latest_model) if latest_model else config.MODEL_NAME
        logger.info(f"기반 모델: {model_path}")

        processor = LayoutLMv3Processor.from_pretrained(
            latest_model and str(latest_model) or config.MODEL_NAME,
            apply_ocr=False
        )
        model = LayoutLMv3ForSequenceClassification.from_pretrained(
            model_path,
            num_labels=len(label2id),
            id2label=id2label,
            label2id=label2id,
            ignore_mismatched_sizes=True,
        ).to(config.DEVICE)

        # 전략 자동 선택
        strategy = req.strategy
        if strategy == "auto":
            buf            = load_replay_buffer()
            existing_types = set(buf.keys())
            new_types      = {id2label[s["label"]] for s in samples}
            strategy       = "replay" if (new_types - existing_types) else "ewc"
        train_status.message = f"전략: {strategy}"
        logger.info(f"학습 전략: {strategy}")

        train_s, val_s = train_test_split(
            samples, test_size=0.2, random_state=42,
            stratify=[s["label"] for s in samples],
        )

        if strategy == "replay":
            replay_s = [s for v in load_replay_buffer().values() for s in v]
            combined = train_s + replay_s
            random.shuffle(combined)
            run_train_loop(model, processor, combined, val_s, req, label2id, strategy, save_path)
            update_replay_buffer(train_s, id2label)
        else:
            run_train_loop(model, processor, train_s, val_s, req, label2id, strategy, save_path)
            update_replay_buffer(train_s, id2label)

        train_status.message = f"학습 완료 → {save_path}"
        logger.info(f"학습 완료: {save_path}")

    except Exception:
        train_status.message = f"오류: {traceback.format_exc()}"
        logger.exception("학습 중 예외 발생")
    finally:
        stop_requested = False
        train_status.running = False


# ── 평가 진입점 ───────────────────────────────────
def run_evaluate(test_size: float = 0.2) -> dict:
    samples  = load_cache()
    label2id = load_label2id()
    id2label = {v: k for k, v in label2id.items()}

    _, test_s = train_test_split(
        samples, test_size=test_size, random_state=42,
        stratify=[s["label"] for s in samples],
    )

    model_path = _get_latest_model_path()
    if not model_path:
        raise RuntimeError("저장된 모델이 없습니다.")

    processor = LayoutLMv3Processor.from_pretrained(
        str(model_path), apply_ocr=False
    )
    model = LayoutLMv3ForSequenceClassification.from_pretrained(
        str(model_path)
    ).to(config.DEVICE)
    model.eval()

    loader      = DataLoader(InvoiceDataset(test_s, processor), batch_size=config.BATCH_SIZE)
    all_preds, all_labels = [], []
    with torch.no_grad():
        for batch in loader:
            out = model(
                input_ids      = batch["input_ids"].to(config.DEVICE),
                attention_mask = batch["attention_mask"].to(config.DEVICE),
                bbox           = batch["bbox"].to(config.DEVICE),
                pixel_values   = batch["pixel_values"].to(config.DEVICE),
            )
            all_preds.extend(out.logits.argmax(dim=-1).cpu().tolist())
            all_labels.extend(batch["labels"].tolist())

    target_names = [id2label[i] for i in range(len(id2label))]
    report_dict  = classification_report(
        all_labels, all_preds, target_names=target_names, output_dict=True
    )
    report_str   = classification_report(
        all_labels, all_preds, target_names=target_names
    )
    logger.info(f"평가 완료 | accuracy={report_dict['accuracy']:.4f}")
    return {
        "test_samples": len(test_s),
        "accuracy":     round(report_dict["accuracy"], 4),
        "report":       report_dict,
        "report_text":  report_str,
    }