package com.easyagents.agent.runtime.event.interceptor;

import com.easyagents.agent.runtime.AgentRuntimeException;
import com.easyagents.agent.runtime.agentscope.AgentScopeMessageAdapter;
import com.easyagents.agent.runtime.event.AgentRuntimeInterceptor;
import com.easyagents.agent.runtime.media.AgentMediaResolver;
import com.easyagents.agent.runtime.media.AgentMediaResource;
import io.agentscope.core.hook.HookEvent;
import io.agentscope.core.hook.PreReasoningEvent;
import io.agentscope.core.message.*;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

/**
 * 在推理前将持久化的业务媒体引用解析为模型可用的 Base64 内容。
 *
 * <p>输入消息会被深拷贝，AgentScope memory/session 中仍保留小体积稳定引用。</p>
 */
public class MediaReferenceInterceptor implements AgentRuntimeInterceptor {

    private final AgentMediaResolver mediaResolver;

    /**
     * 创建媒体引用干预器。
     *
     * @param mediaResolver 媒体引用解析器
     */
    public MediaReferenceInterceptor(AgentMediaResolver mediaResolver) {
        this.mediaResolver = mediaResolver;
    }

    /**
     * 解析推理输入中的内部媒体引用。
     *
     * @param event AgentScope Hook 事件
     * @param <T> Hook 事件类型
     * @return 处理后的事件
     */
    @Override
    public <T extends HookEvent> Mono<T> intercept(T event) {
        if (event instanceof PreReasoningEvent preReasoningEvent) {
            preReasoningEvent.setInputMessages(resolveMessages(preReasoningEvent.getInputMessages()));
        }
        return Mono.just(event);
    }

    /**
     * 在 AutoContext 完成消息重写后执行媒体解析。
     *
     * @return 执行优先级
     */
    @Override
    public int priority() {
        return 10;
    }

    private List<Msg> resolveMessages(List<Msg> messages) {
        if (messages == null || messages.isEmpty()) {
            return List.of();
        }
        List<Msg> resolved = new ArrayList<>(messages.size());
        for (Msg message : messages) {
            List<ContentBlock> content = resolveBlocks(message.getContent());
            resolved.add(Msg.builder()
                    .id(message.getId())
                    .name(message.getName())
                    .role(message.getRole())
                    .content(content)
                    .metadata(message.getMetadata())
                    .timestamp(message.getTimestamp())
                    .build());
        }
        return resolved;
    }

    private List<ContentBlock> resolveBlocks(List<ContentBlock> blocks) {
        if (blocks == null || blocks.isEmpty()) {
            return List.of();
        }
        List<ContentBlock> resolved = new ArrayList<>(blocks.size());
        for (ContentBlock block : blocks) {
            if (block instanceof ImageBlock imageBlock) {
                resolved.add(resolveImage(imageBlock));
            } else if (block instanceof AudioBlock audioBlock) {
                resolved.add(resolveAudio(audioBlock));
            } else if (block instanceof VideoBlock videoBlock) {
                resolved.add(resolveVideo(videoBlock));
            } else {
                resolved.add(block);
            }
        }
        return resolved;
    }

    private ImageBlock resolveImage(ImageBlock block) {
        AgentMediaResource resource = resolve(block.getSource());
        if (resource == null) {
            return block;
        }
        return ImageBlock.builder()
                .source(base64Source(resource))
                .minPixels(block.getMinPixels())
                .maxPixels(block.getMaxPixels())
                .build();
    }

    private AudioBlock resolveAudio(AudioBlock block) {
        AgentMediaResource resource = resolve(block.getSource());
        return resource == null ? block : AudioBlock.builder().source(base64Source(resource)).build();
    }

    private VideoBlock resolveVideo(VideoBlock block) {
        AgentMediaResource resource = resolve(block.getSource());
        if (resource == null) {
            return block;
        }
        return VideoBlock.builder()
                .source(base64Source(resource))
                .fps(block.getFps())
                .maxFrames(block.getMaxFrames())
                .minPixels(block.getMinPixels())
                .maxPixels(block.getMaxPixels())
                .totalPixels(block.getTotalPixels())
                .build();
    }

    private AgentMediaResource resolve(Source source) {
        if (!(source instanceof URLSource urlSource)
                || urlSource.getUrl() == null
                || !urlSource.getUrl().startsWith(AgentScopeMessageAdapter.MEDIA_REFERENCE_SCHEME)) {
            return null;
        }
        if (mediaResolver == null) {
            throw new AgentRuntimeException("Agent media resolver is required for media references.");
        }
        String reference;
        try {
            reference = AgentScopeMessageAdapter.decodeReference(urlSource.getUrl());
        } catch (IllegalArgumentException error) {
            throw new AgentRuntimeException("Agent media reference is invalid.", error);
        }
        AgentMediaResource resource = mediaResolver.resolve(reference);
        if (resource == null || resource.bytes().length == 0 || resource.mimeType() == null
                || resource.mimeType().isBlank()) {
            throw new AgentRuntimeException("Agent media resource is empty.");
        }
        return resource;
    }

    private Base64Source base64Source(AgentMediaResource resource) {
        return Base64Source.builder()
                .mediaType(resource.mimeType())
                .data(Base64.getEncoder().encodeToString(resource.bytes()))
                .build();
    }
}
