# This file is the implementation of custom agent structured agent block florent-block
from dataiku.llm.python.blocks_graph import BlockHandler, NextBlock
import json
import requests
import logging

logger = logging.getLogger(__name__)


class CustomStructuredAgentBlock(BlockHandler):
    def __init__(self, turn, sequence_context, block_config):
        super().__init__(turn, sequence_context, block_config)
        self.story_id = self.turn.state_get("story_id")
        if not self.story_id:
            raise ValueError("Missing 'story_id' in state")

        preset_value = self.block_config.get("config").get("api_key")
        if not preset_value:
            raise ValueError("Missing preset 'api_key' in block config")
        self.api_token = preset_value.get("shortcut_api_token")
        if not self.api_token:
            raise ValueError("Missing 'shortcut-token' in preset 'api_token'")

        self.send_to_user = self.block_config.get("config").get("send_to_user")
        self.state_key = self.block_config.get("config").get("state_key")

    def process_stream(self, trace):
        url = f"https://api.app.shortcut.com/api/v3/stories/{self.story_id}"
        headers = {
            "Content-Type": "application/json",
            "Shortcut-Token": self.api_token,
        }

        response = requests.get(url, headers=headers, timeout=30)
        if response.status_code == 404:
            raise ValueError(f"Story {self.story_id} not found")
        else:
            response.raise_for_status()

            story = response.json()

            self.turn.state_set(self.state_key, story)

            trace.outputs["output"] = {
                "story_id": self.story_id,
                "story_name": story.get("name"),
                "app_url": story.get("app_url"),
            }

            logger.info(json.dumps(story, indent=2))
            if (self.send_to_user):
                yield {
                    "chunk": {
                        "text": f"""Shortcut story {self.story_id} fetched: {story.get('name', '<no name>')}
                        {story}"""
                    }
                }
            yield NextBlock(id=self.block_config.get("defaultNextBlock"))