from typing import List

from pygeai.core.omni.responses import OmniParserResponse, OmniParserPart, OmniParserUsage


class OmniParserResponseMapper:

    @classmethod
    def map_to_omni_parser_response(cls, data: dict) -> OmniParserResponse:
        parts_list = cls.map_to_part_list(data.get('parts', []))
        usage = cls.map_to_usage(data.get('usage'))

        return OmniParserResponse(
            status=data.get('status', ''),
            parts=parts_list,
            request_id=data.get('requestId', ''),
            usage=usage
        )

    @classmethod
    def map_to_part_list(cls, data: list) -> List[OmniParserPart]:
        part_list = []
        for item in data:
            part = cls.map_to_part(item)
            part_list.append(part)
        return part_list

    @classmethod
    def map_to_part(cls, data: dict) -> OmniParserPart:
        return OmniParserPart(
            page=data.get('page'),
            type=data.get('type'),
            text=data.get('text', ''),
            mid_time=data.get('midTime')
        )

    @classmethod
    def map_to_usage(cls, data: dict) -> OmniParserUsage | None:
        if data is None:
            return None

        return OmniParserUsage(
            total_cost=data.get('total_cost'),
            total_tokens=data.get('total_tokens'),
            completion_tokens=data.get('completion_tokens'),
            prompt_tokens=data.get('prompt_tokens'),
            prompt_cost=data.get('prompt_cost'),
            completion_cost=data.get('completion_cost'),
            currency=data.get('currency'),
            prompt_tokens_details=data.get('prompt_tokens_details'),
            completion_tokens_details=data.get('completion_tokens_details')
        )
