import type {
  Cluster, Campus, DataCenter, DeploymentCohort,
  DCNetworkSummary, CampusNetworkSummary, ClusterNetworkSummary,
  TrainingPipeline, BaseModel, ModelFamily,
} from '@ai-tycoon/shared';
import { uuid } from '@ai-tycoon/shared';
import type { DeepPartial } from './createTestState';

function emptyDCNetwork(): DCNetworkSummary {
  return {
    totalByTier: {},
    healthyByTier: {},
    repairBatches: [],
    networkRackCount: 0,
    racksDisconnected: 0,
    racksDegraded: 0,
    averageBandwidth: 1,
    effectiveFlopsFraction: 1,
  };
}

function emptyCampusNetwork(): CampusNetworkSummary {
  return { totalT4: 0, healthyT4: 0, crossDCBandwidth: 1 };
}

function emptyClusterNetwork(): ClusterNetworkSummary {
  return { totalT5: 0, healthyT5: 0, crossCampusBandwidth: 1 };
}

export function createTestDataCenter(overrides?: DeepPartial<DataCenter>): DataCenter {
  const base: DataCenter = {
    id: uuid(),
    name: 'Test DC',
    campusId: '',
    tier: 'small',
    status: 'operational',
    constructionProgress: 0,
    constructionTotal: 0,
    rackSkuId: 't4-x4',
    computeRacksOnline: 4,
    computeRacksFailed: 0,
    networkSummary: emptyDCNetwork(),
    deploymentCohorts: [],
    retrofitState: null,
    coolingLevel: 0,
    redundancyLevel: 0,
    coolingType: 'air',
    networkFabric: 'ethernet-100g',
    effectiveComputeRacks: 4,
    usedSlots: 4,
    usedPowerKW: 20,
    energyCostPerTick: 5,
    maintenanceCostPerTick: 2,
    currentUptime: 1,
    dcTrainingFlops: 1e12,
    dcInferenceFlops: 1e12,
    dcTotalVramGB: 64,
  };
  return overrides ? { ...base, ...overrides } as DataCenter : base;
}

export function createTestCampus(overrides?: DeepPartial<Campus>): Campus {
  const dc = createTestDataCenter();
  const campusId = uuid();
  dc.campusId = campusId;
  const base: Campus = {
    id: campusId,
    name: 'Test Campus',
    clusterId: '',
    dcTier: 'small',
    dataCenters: [dc],
    status: 'operational',
    constructionProgress: 0,
    constructionTotal: 0,
    retrofitQueue: null,
    networkSummary: emptyCampusNetwork(),
  };
  return overrides ? { ...base, ...overrides } as Campus : base;
}

export function createTestCluster(overrides?: DeepPartial<Cluster>): Cluster {
  const campus = createTestCampus();
  const clusterId = uuid();
  campus.clusterId = clusterId;
  campus.dataCenters[0].campusId = campus.id;
  const base: Cluster = {
    id: clusterId,
    name: 'Test Cluster',
    locationId: 'us-east',
    campuses: [campus],
    status: 'operational',
    constructionProgress: 0,
    constructionTotal: 0,
    networkSummary: emptyClusterNetwork(),
  };
  return overrides ? { ...base, ...overrides } as Cluster : base;
}

export function createTestTrainingPipeline(overrides?: DeepPartial<TrainingPipeline>): TrainingPipeline {
  const base: TrainingPipeline = {
    id: uuid(),
    familyId: 'test-family',
    modelName: 'Test Model',
    architecture: {
      type: 'dense',
      totalParameters: 7e9,
      activeParameters: 7e9,
      contextWindow: 8192,
      vocabularySize: 32000,
    },
    dataMix: { web: 0.4, code: 0.2, books: 0.15, academic: 0.1, conversational: 0.1, specialized: 0.05 },
    currentStage: 'pretraining',
    stages: {
      pretraining: {
        targetTokens: 1e12,
        processedTokens: 0,
        computeAllocated: 0,
        progressTicks: 0,
        totalTicks: 1000,
        lossValue: 4.0,
        chinchillaRatio: 1.0,
        isComplete: false,
      },
      sft: {
        specializations: ['general'],
        progressTicks: 0,
        totalTicks: 100,
        isComplete: false,
      },
      alignment: {
        method: 'rlhf',
        safetyWeight: 0.5,
        helpfulnessWeight: 0.5,
        progressTicks: 0,
        totalTicks: 80,
        isComplete: false,
      },
    },
    status: 'active',
    allocatedComputeFraction: 1.0,
    events: [],
    startedAtTick: 0,
    sizeTier: 'small',
    isPointRelease: false,
    sourceModelId: null,
  };
  return overrides ? { ...base, ...overrides } as TrainingPipeline : base;
}

export function createTestBaseModel(overrides?: Partial<BaseModel>): BaseModel {
  const base: BaseModel = {
    id: uuid(),
    familyId: 'test-family',
    name: 'Test Model v1',
    architecture: {
      type: 'dense',
      totalParameters: 7e9,
      activeParameters: 7e9,
      contextWindow: 8192,
      vocabularySize: 32000,
    },
    rawCapability: 40,
    capabilityScore: 40,
    safetyScore: 50,
    qualityScore: 40,
    sftSpecializations: ['general'],
    alignmentMethod: 'rlhf',
    completedAtTick: 100,
    isDeployed: true,
    isOpenSourced: false,
    sizeTier: 'small',
    isPointRelease: false,
    sourceModelId: null,
    dataMix: { web: 0.4, code: 0.2, books: 0.15, academic: 0.1, conversational: 0.1, specialized: 0.05 },
  };
  return overrides ? { ...base, ...overrides } : base;
}

export function createTestModelFamily(overrides?: Partial<ModelFamily>): ModelFamily {
  const base: ModelFamily = {
    id: uuid(),
    name: 'Test Family',
    generation: 1,
    baseModelIds: [],
    variants: [],
    createdAtTick: 0,
  };
  return overrides ? { ...base, ...overrides } : base;
}

export function createTestDeploymentCohort(overrides?: Partial<DeploymentCohort>): DeploymentCohort {
  const base: DeploymentCohort = {
    id: uuid(),
    count: 4,
    skuId: 't4-x4',
    stage: 'production' as any,
    stageProgress: 0,
    stageTotal: 0,
    repairCount: 0,
  };
  return overrides ? { ...base, ...overrides } : base;
}