#!/usr/bin/env python
"""
数据迁移脚本：从 lobsters 表迁移到 agents 表

使用方法：
python manage.py migrate_data
"""

from django.core.management.base import BaseCommand
from lobsters.models import Lobster, LobsterDiary
from agents.models import Agent, AgentDiary


class Command(BaseCommand):
    help = '从 lobsters 表迁移数据到 agents 表'

    def handle(self, *args, **kwargs):
        self.stdout.write('🚀 开始数据迁移...\n')
        
        # 迁移 Lobster → Agent
        self.stdout.write('📊 迁移 Lobster 数据...')
        lobsters = Lobster.objects.all()
        agent_count = 0
        
        for lobster in lobsters:
            Agent.objects.update_or_create(
                id=lobster.id,
                defaults={
                    'name': lobster.name,
                    'emoji': lobster.emoji,
                    'port': lobster.port,
                    'specialty': lobster.specialty,
                    'container': lobster.container,
                    'app_name': lobster.app_name,
                    'app_id': lobster.app_id,
                    'workspace': lobster.workspace,
                }
            )
            agent_count += 1
        
        self.stdout.write(self.style.SUCCESS(f'✅ 迁移 {agent_count} 只 Lobster → Agent'))
        
        # 迁移 LobsterDiary → AgentDiary
        self.stdout.write('\n📝 迁移 Diary 数据...')
        diaries = LobsterDiary.objects.all()
        diary_count = 0
        
        for diary in diaries:
            AgentDiary.objects.update_or_create(
                id=diary.id,
                defaults={
                    'agent_id': diary.lobster_id,
                    'date': diary.date,
                    'title': diary.title,
                    'content': diary.content,
                    'category': diary.category,
                    'tags': diary.tags,
                    'embedding': diary.embedding,
                    'embedding_model': diary.embedding_model,
                }
            )
            diary_count += 1
        
        self.stdout.write(self.style.SUCCESS(f'✅ 迁移 {diary_count} 篇 Diary'))
        
        # 统计
        self.stdout.write(self.style.SUCCESS(f'\n🎉 迁移完成！'))
        self.stdout.write(f'   Agents: {Agent.objects.count()}')
        self.stdout.write(f'   AgentDiaries: {AgentDiary.objects.count()}')