import { ComponentErrorBoundary } from '../components/common/ComponentErrorBoundary';
import { ModelCard } from '../components/ml/ModelCard';
import { EnsemblePanel } from '../components/ml/EnsemblePanel';
import { TrainingProgress } from '../components/ml/TrainingProgress';
import { useMlPredictions, useTrainingJobs } from '../hooks/useApi';
import { useMetricsStream } from '../hooks/useGrpcStream';

const MODEL_GROUPS = [
  { label: 'Reinforcement Learning', models: ['DQN', 'PPO'] },
  { label: 'Temporal', models: ['TFT', 'Mamba2', 'xLSTM'] },
  { label: 'Graph / Structure', models: ['TGGN', 'TLOB', 'LNN'] },
  { label: 'Generative', models: ['KAN', 'Diffusion'] },
] as const;

export function MLDashboard() {
  const predictions = useMlPredictions();
  const trainingJobs = useTrainingJobs();
  const { messages: metricsMessages } = useMetricsStream();

  // Extract regime data from most recent metrics stream message
  const latestMetrics = metricsMessages.length > 0
    ? metricsMessages.at(-1)
    : undefined;

  // The gRPC response wraps predictions in a `predictions` field
  const preds = predictions.data?.predictions ?? [];

  // Adapt proto MLPrediction to the shape the ModelCard / EnsemblePanel expects
  const adaptedPreds = preds.map((p) => {
    const raw = p.ensembleAction?.toLowerCase();
    const signal: 'buy' | 'sell' | 'hold' =
      raw === 'buy' || raw === 'sell' ? raw : 'hold';
    return {
      model: p.ensembleAction,
      symbol: p.symbol,
      signal,
      confidence: p.ensembleConfidence,
      predicted_return: p.ensembleSignal,
      timestamp: '',
    };
  });

  return (
    <div className="space-y-4">
      {(predictions.isError || trainingJobs.isError) && (
        <div className="px-3 py-2 rounded border border-[var(--color-red)] bg-red-500/10 text-sm text-[var(--color-red)] flex items-center justify-between">
          <span>Failed to load {predictions.isError ? 'predictions' : 'training jobs'}: {(predictions.error ?? trainingJobs.error)?.message}</span>
          <button onClick={() => { predictions.refetch(); trainingJobs.refetch(); }} className="text-xs underline">Retry</button>
        </div>
      )}

      {/* Model Cards grouped by category */}
      <ComponentErrorBoundary name="Model Cards">
        <div className="space-y-4">
          {MODEL_GROUPS.map((group) => (
            <div key={group.label}>
              <h3 className="text-xs font-semibold uppercase tracking-wider text-[var(--color-text-secondary)] mb-2">
                {group.label}
              </h3>
              <div className="grid grid-cols-1 sm:grid-cols-2 xl:grid-cols-5 gap-4">
                {group.models.map((model) => {
                  const pred = adaptedPreds.find((p) => p.model === model);
                  return (
                    <div
                      key={model}
                      className="rounded border border-[var(--color-border)] bg-[var(--color-bg-card)]"
                    >
                      <ModelCard
                        model={model}
                        signal={pred?.signal}
                        confidence={pred?.confidence}
                        predictedReturn={pred?.predicted_return}
                      />
                    </div>
                  );
                })}
              </div>
            </div>
          ))}
        </div>
      </ComponentErrorBoundary>

      {/* Ensemble + Regime */}
      <div className="grid grid-cols-1 md:grid-cols-2 gap-4">
        <ComponentErrorBoundary name="Ensemble Panel">
          <div className="rounded border border-[var(--color-border)] bg-[var(--color-bg-card)]">
            <EnsemblePanel predictions={adaptedPreds} />
          </div>
        </ComponentErrorBoundary>
        <div className="rounded border border-[var(--color-border)] bg-[var(--color-bg-card)]">
          <div className="px-3 py-1.5 border-b border-[var(--color-border)]">
            <span className="text-sm font-medium">Market Regime</span>
          </div>
          <div className="p-3 space-y-2">
            <div className="flex justify-between text-xs">
              <span className="text-[var(--color-text-secondary)]">Current Regime</span>
              <span className="text-[var(--color-accent)]">--</span>
            </div>
            <div className="flex justify-between text-xs">
              <span className="text-[var(--color-text-secondary)]">Volatility</span>
              <span>--</span>
            </div>
            <div className="flex justify-between text-xs">
              <span className="text-[var(--color-text-secondary)]">Regime Confidence</span>
              <span>--</span>
            </div>
            <div className="flex justify-between text-xs">
              <span className="text-[var(--color-text-secondary)]">Regime Duration</span>
              <span>{latestMetrics ? 'streaming' : '--'}</span>
            </div>
          </div>
        </div>
      </div>

      {/* Training Progress */}
      <ComponentErrorBoundary name="Training Progress">
        <div className="rounded border border-[var(--color-border)] bg-[var(--color-bg-card)]">
          <TrainingProgress
            jobs={trainingJobs.data?.jobs}
            loading={trainingJobs.isLoading}
          />
        </div>
      </ComponentErrorBoundary>
    </div>
  );
}
