import asyncio
from .base import BaseAgent
from ..db import create_task, update_task_status, approve_task, reject_task, add_log
from .. import service_proxy
from .. import telegram_bot

class MusicAgent(BaseAgent):
    agent_id = "music"
    display_name = "음악 프로듀서"

    async def on_schedule(self) -> None:
        pass

    async def on_command(self, command: str, params: dict) -> dict:
        if command == "compose":
            prompt = params.get("prompt", "")
            style = params.get("style", "")
            model = params.get("model", "V4")
            instrumental = params.get("instrumental", False)

            if not prompt:
                return {"ok": False, "message": "프롬프트를 입력해주세요"}

            task_id = create_task(self.agent_id, "compose", {
                "prompt": prompt, "style": style,
                "model": model, "instrumental": instrumental,
            }, requires_approval=True)

            await self.transition("waiting", "프롬프트 승인 대기", task_id)

            detail = f"프롬프트: {prompt}"
            if style:
                detail += f"\n스타일: {style}"
            detail += f"\n모델: {model}"

            await telegram_bot.send_approval_request(
                self.agent_id, task_id,
                "🎵 [음악 에이전트] 작곡 요청", detail,
            )

            return {"ok": True, "task_id": task_id, "message": "승인 대기 중"}

        if command == "credits":
            credits = await service_proxy.get_music_credits()
            return {"ok": True, "credits": credits}

        return {"ok": False, "message": f"Unknown command: {command}"}

    async def on_approval(self, task_id: str, approved: bool, feedback: str = "") -> None:
        if not approved:
            reject_task(task_id)
            await self.transition("idle", "작곡 거절됨")
            await telegram_bot.send_task_result(
                self.agent_id, "🎵 [음악 에이전트] 작곡 취소",
                "사용자가 거절했습니다.",
            )
            return

        from ..db import get_task
        task = get_task(task_id)
        if not task:
            return

        approve_task(task_id, via="telegram")
        await self.transition("working", "작곡 중...", task_id)
        asyncio.create_task(self._poll_composition(task_id, task))

    async def _poll_composition(self, task_id: str, task: dict) -> None:
        try:
            input_data = task["input_data"]
            payload = {
                "provider": "suno",
                "model": input_data.get("model", "V4"),
                "prompt": input_data.get("prompt", ""),
                "style": input_data.get("style", ""),
                "instrumental": input_data.get("instrumental", False),
                "custom_mode": True,
            }

            result = await service_proxy.generate_music(payload)
            music_task_id = result.get("task_id")

            if not music_task_id:
                raise Exception("music-lab did not return task_id")

            for _ in range(60):
                await asyncio.sleep(5)
                status = await service_proxy.get_music_status(music_task_id)
                state = status.get("status", "")

                if state == "succeeded":
                    tracks = status.get("tracks", [])
                    update_task_status(task_id, "succeeded", {
                        "music_task_id": music_task_id,
                        "tracks": tracks,
                    })
                    await self.transition("reporting", "작곡 완료!")

                    track_info = ""
                    for t in tracks:
                        title = t.get("title", "Untitled")
                        url = t.get("audio_url", "")
                        track_info += f"🎶 {title}\n{url}\n"

                    await telegram_bot.send_task_result(
                        self.agent_id, "🎵 [음악 에이전트] 작곡 완료",
                        track_info or "트랙 생성 완료",
                    )
                    await self.transition("idle", "작곡 완료")
                    return

                if state == "failed":
                    raise Exception(status.get("message", "Generation failed"))

            raise Exception("Timeout: 5분 초과")

        except Exception as e:
            add_log(self.agent_id, f"Compose failed: {e}", "error", task_id)
            update_task_status(task_id, "failed", {"error": str(e)})
            await self.transition("idle", f"오류: {e}")
            await telegram_bot.send_task_result(
                self.agent_id, "🎵 [음악 에이전트] 작곡 실패",
                f"오류: {e}",
            )