#!/usr/bin/env bun

const API_KEY = process.env.API_KEY;
const PROJECT_ID = 'your-project-id-here';
const FINETUNED_MODEL_ALIAS = 'your-model-alias';
const JUDGE_MODEL = 'claude-4.5-sonnet';

// CRITICAL: Replace with prompts matching your model's domain!
const TEST_PROMPTS = [
	'YOUR_FIRST_TEST_PROMPT_HERE',
	'YOUR_SECOND_TEST_PROMPT_HERE',
	'YOUR_THIRD_TEST_PROMPT_HERE'
];

if (!API_KEY) {
	console.error('Error: API_KEY environment variable is required');
	process.exit(1);
}

function sleep(ms: number) {
	return new Promise((r) => setTimeout(r, ms));
}

async function main() {
	const projectRes = await fetch(`https://studio.premai.io/api/v1/public/projects/${PROJECT_ID}`, {
		headers: { 'Authorization': `Bearer ${API_KEY}` }
	});
	if (!projectRes.ok) throw new Error(`Failed to fetch project: ${projectRes.status}`);
	const project = await projectRes.json();
	const dataset = project.project.children.find((child: any) => child.type === 'dataset');
	if (!dataset) throw new Error('No dataset found in project');
	const DATASET_ID = dataset.id;

	const modelResponses = [];
	for (const prompt of TEST_PROMPTS) {
		const res = await fetch('https://studio.premai.io/api/v1/chat/completions', {
			method: 'POST',
			headers: {
				'Authorization': `Bearer ${API_KEY}`,
				'Content-Type': 'application/json'
			},
			body: JSON.stringify({
				project_id: PROJECT_ID,
				model: FINETUNED_MODEL_ALIAS,
				messages: [{ role: 'user', content: prompt }]
			})
		});
		if (!res.ok) throw new Error(`${res.status}: ${await res.text()}`);
		const response = await res.json();
		modelResponses.push({ prompt, answer: response.choices[0].message.content });
	}

	const evaluations = [];

	for (const item of modelResponses) {
		const judgePrompt = `You are an expert AI evaluator. Evaluate the following response.

User Question: ${item.prompt}

Model Response: ${item.answer}

Provide your evaluation in the following JSON format (output ONLY the JSON, no other text):
{
  "feedback": "<detailed explanation highlighting strengths and weaknesses>",
  "reasoning": "<why you gave this specific score>",
  "score": <number between 0 and 1, where 0 is completely wrong and 1 is perfect>
}`;

		const res = await fetch('https://studio.premai.io/api/v1/chat/completions', {
			method: 'POST',
			headers: {
				'Authorization': `Bearer ${API_KEY}`,
				'Content-Type': 'application/json'
			},
			body: JSON.stringify({
				project_id: PROJECT_ID,
				model: JUDGE_MODEL,
				messages: [{ role: 'user', content: judgePrompt }],
				temperature: 0.1
			})
		});

		if (!res.ok) throw new Error(`${res.status}: ${await res.text()}`);
		const response = await res.json();
		const judgeResponse = response.choices[0].message.content;

		let evaluation;
		try {
			const jsonMatch = judgeResponse.match(/\{[\s\S]*\}/);
			evaluation = JSON.parse(jsonMatch ? jsonMatch[0] : judgeResponse);
		} catch (e) {
			console.warn(`Warning: Could not parse judge response`);
			evaluation = {
				score: 5,
				feedback: judgeResponse,
				reasoning: 'Could not parse structured evaluation'
			};
		}

		evaluations.push({
			prompt: item.prompt,
			answer: item.answer,
			score: evaluation.score,
			feedback: evaluation.feedback,
			reasoning: evaluation.reasoning
		});

		console.log(`✓ Evaluated: "${item.prompt.substring(0, 40)}..." - Score: ${evaluation.score}`);
	}

	console.log(`\n✓ Evaluated ${evaluations.length} responses\n`);

	// Step 3: Create traces
	console.log('=== Step 3: Creating traces with evaluation data ===\n');

	const traceIds = [];

	for (const evaluation of evaluations) {
		const res = await fetch('https://studio.premai.io/api/v1/traces', {
			method: 'POST',
			headers: {
				'Authorization': `Bearer ${API_KEY}`,
				'Content-Type': 'application/json'
			},
			body: JSON.stringify({
				project_id: PROJECT_ID,
				model_id: FINETUNED_MODEL_ALIAS,
				input: evaluation.prompt,
				output: evaluation.answer,
				score: evaluation.score,
				feedback: `${evaluation.feedback}\n\nReasoning: ${evaluation.reasoning}`
			})
		});

		if (!res.ok) throw new Error(`${res.status}: ${await res.text()}`);
		const trace = await res.json();

		traceIds.push(trace.id);
		console.log(`✓ Created trace ${trace.id} - Score: ${evaluation.score}`);
	}

	console.log(`\n✓ Created ${traceIds.length} traces\n`);

	// Step 4: Add ALL traces to dataset (including low-scoring ones!)
	console.log('=== Step 4: Adding ALL traces to dataset ===\n');
	console.log('IMPORTANT: Adding ALL traces, including low-scoring ones!');
	console.log('Low-scoring traces help identify model weaknesses.');
	console.log('The addToDataset endpoint automatically uses the dataset from the project.\n');

	const addedTraces = [];

	for (let i = 0; i < traceIds.length; i++) {
		const traceId = traceIds[i];
		const score = evaluations[i].score;

		const res = await fetch(`https://studio.premai.io/api/v1/traces/${traceId}/addToDataset`, {
			method: 'POST',
			headers: {
				'Authorization': `Bearer ${API_KEY}`
			}
		});

		if (!res.ok) {
			console.warn(`Warning: Failed to add trace ${traceId}`);
			continue;
		}

		addedTraces.push(traceId);
		const quality = score >= 0.7 ? 'high-quality' : 'low-quality (will be improved)';
		console.log(`✓ Added trace ${traceId} (${quality}, score: ${score}) to dataset`);
	}

	console.log(`\n✓ Added ${addedTraces.length} traces to dataset\n`);

	// Step 5: Create new snapshot
	console.log('=== Step 5: Creating new snapshot ===\n');

	const res5 = await fetch('https://studio.premai.io/api/v1/public/snapshots/create', {
		method: 'POST',
		headers: {
			'Authorization': `Bearer ${API_KEY}`,
			'Content-Type': 'application/json'
		},
		body: JSON.stringify({
			dataset_id: DATASET_ID,
			split_percentage: 80
		})
	});

	if (!res5.ok) throw new Error(`${res5.status}: ${await res5.text()}`);
	const { snapshot_id } = await res5.json();

	console.log(`✓ Created new snapshot: ${snapshot_id}\n`);

	// Step 6: Generate recommendations
	console.log('=== Step 6: Generating recommendations ===\n');

	const res6 = await fetch('https://studio.premai.io/api/v1/public/recommendations/generate', {
		method: 'POST',
		headers: {
			'Authorization': `Bearer ${API_KEY}`,
			'Content-Type': 'application/json'
		},
		body: JSON.stringify({ snapshot_id })
	});

	if (!res6.ok) throw new Error(`${res6.status}: ${await res6.text()}`);

	let recs;
	do {
		await sleep(5000);
		const res = await fetch(`https://studio.premai.io/api/v1/public/recommendations/${snapshot_id}`, {
			headers: { 'Authorization': `Bearer ${API_KEY}` }
		});
		if (!res.ok) throw new Error(`${res.status}: ${await res.text()}`);
		recs = await res.json();
	} while (recs.status === 'processing');

	console.log('✓ Recommendations ready\n');

	// Step 7: Launch new fine-tuning job
	console.log('=== Step 7: Launching new fine-tuning job ===\n');

	const experiments = recs.recommended_experiments
		.filter((e: any) => e.recommended)
		.map(({ recommended, reason_for_recommendation, ...experiment }: any) => experiment);

	if (experiments.length === 0) {
		console.error('✗ No recommended experiments found');
		process.exit(1);
	}

	const res7 = await fetch('https://studio.premai.io/api/v1/public/finetuning/create', {
		method: 'POST',
		headers: {
			'Authorization': `Bearer ${API_KEY}`,
			'Content-Type': 'application/json'
		},
		body: JSON.stringify({
			snapshot_id,
			name: `Continuous Fine-tuning - ${new Date().toISOString().split('T')[0]}`,
			experiments
		})
	});

	if (!res7.ok) throw new Error(`${res7.status}: ${await res7.text()}`);
	const { job_id } = await res7.json();

	console.log(`✓ Fine-tuning job started: ${job_id}\n`);
	console.log('✓ Continuous fine-tuning cycle complete!\n');
}

main().catch((err) => {
	console.error('\n✗ Error:', err.message);
	process.exit(1);
});