from DbKit.Db import Db
from datetime import datetime

class DouYinModel:
    def __init__(self, db: Db = None):
        self.db = db or Db()

    async def init(self):
        await self.db.init_db()

    async def update_status(self, id, status, error_msg=None):
        await self.db.init_db()
        params = {
            "id": id,
            "status": status,
            "error_msg": error_msg
        }
        return await self.db.execute_update("DouYin.updateStatus", params)

    async def update_record(self, id, title, obs_url, transcript, status):
        await self.db.init_db()
        # Truncate title to 100 chars to fit DB schema
        if title and len(title) > 100:
            title = title[:100] + "..."
            
        params = {
            "id": id,
            "title": title,
            "obs_url": obs_url,
            "transcript": transcript,
            "status": status
        }
        return await self.db.execute_update("DouYin.updateRecord", params)

    async def insert_record(self, id, url, create_time=None):
        await self.db.init_db()
        if create_time is None:
            create_time = datetime.now()
        params = {
            "id": id,
            "url": url,
            "create_time": create_time
        }
        return await self.db.execute_update("DouYin.insertRecord", params)

    async def get_records(self, limit=50):
        await self.db.init_db()
        params = {"limit": limit}
        records = await self.db.find("DouYin.getRecords", params)
        
        # Manually handle datetime serialization
        for r in records:
            if 'create_time' in r and r['create_time']:
                r['create_time'] = r['create_time'].strftime("%Y-%m-%d %H:%M:%S")
            if 'update_time' in r and r['update_time']:
                r['update_time'] = r['update_time'].strftime("%Y-%m-%d %H:%M:%S")
        return records

    async def delete_record(self, id):
        await self.db.init_db()
        params = {"id": id}
        return await self.db.execute_update("DouYin.deleteRecord", params)

    async def get_transcripts(self, ids=None, limit=20):
        await self.db.init_db()
        if ids:
            params = {"ids": ids}
            return await self.db.find("DouYin.getTranscriptsByIds", params)
        else:
            params = {"limit": limit}
            return await self.db.find("DouYin.getLatestTranscripts", params)

    async def get_interrupted_tasks(self, limit=20):
        await self.db.init_db()
        params = {"limit": limit}
        return await self.db.find("DouYin.getInterruptedTasks", params)

    async def close(self):
        await self.db.shutdown()