modcomms/frontend/services/geminiService.ts

import type { AgentReview, SubReview, AgentName } from '../types';
import { IPublicClientApplication } from '@azure/msal-browser';
import { getAccessToken } from './authService';

// WebSocket URL for backend communication
const WS_URL = import.meta.env.VITE_BACKEND_WS_URL || 'ws://localhost:8000/ws/analyze';
const HTTP_URL = import.meta.env.VITE_BACKEND_URL || 'http://localhost:8000';

/**
 * Analyze a proof using the backend WebSocket API.
 * Provides real-time updates as each agent completes.
 * Now requires MSAL instance to acquire access token.
 */
export const analyzeProof = async (
    file: File,
    onAgentUpdate: (name: AgentName | 'Summary', review?: SubReview) => void,
    msalInstance: IPublicClientApplication
): Promise<AgentReview> => {
    // Acquire token before connecting
    const accessToken = await getAccessToken(msalInstance);
    if (!accessToken) {
        throw new Error('Failed to acquire access token. Please sign in again.');
    }

    return new Promise((resolve, reject) => {
        const ws = new WebSocket(WS_URL);
        let resolved = false;

        ws.onopen = () => {
            // Convert file to base64 and send
            const reader = new FileReader();
            reader.onloadend = () => {
                const base64Data = (reader.result as string).split(',')[1];
                ws.send(JSON.stringify({
                    type: 'analyze',
                    file_data: base64Data,
                    file_type: file.type,
                    is_wip: false,
                    access_token: accessToken
                }));
            };
            reader.onerror = () => {
                ws.close();
                reject(new Error('Failed to read file'));
            };
            reader.readAsDataURL(file);
        };

        ws.onmessage = (event) => {
            try {
                const message = JSON.parse(event.data);

                switch (message.type) {
                    case 'agent_started':
                        // Agent is starting - can optionally show loading state
                        break;

                    case 'agent_completed':
                        // Agent completed - update UI
                        onAgentUpdate(message.agent_name as AgentName, message.review);
                        break;

                    case 'summary':
                        // Summary ready
                        onAgentUpdate('Summary');
                        break;

                    case 'complete':
                        // Analysis complete - resolve with full result
                        resolved = true;
                        ws.close();
                        resolve(message.result as AgentReview);
                        break;

                    case 'error':
                        // Error occurred
                        resolved = true;
                        ws.close();
                        reject(new Error(message.message || 'Analysis failed'));
                        break;
                }
            } catch (e) {
                console.error('Failed to parse WebSocket message:', e);
            }
        };

        ws.onerror = () => {
            if (!resolved) {
                resolved = true;
                reject(new Error('WebSocket connection error. Is the backend running?'));
            }
        };

        ws.onclose = (event) => {
            if (!resolved && !event.wasClean) {
                resolved = true;
                reject(new Error('WebSocket connection closed unexpectedly'));
            }
        };
    });
};

/**
 * Analyze a work-in-progress proof.
 * Returns a conversational summary for the creative team.
 */
export const analyzeWIPProof = async (
    file: File,
    onAgentUpdate: (name: AgentName | 'Summary', review?: SubReview) => void,
    msalInstance: IPublicClientApplication
): Promise<string> => {
    // Acquire token before connecting
    const accessToken = await getAccessToken(msalInstance);
    if (!accessToken) {
        throw new Error('Failed to acquire access token. Please sign in again.');
    }

    return new Promise((resolve, reject) => {
        const ws = new WebSocket(WS_URL);
        let resolved = false;

        ws.onopen = () => {
            const reader = new FileReader();
            reader.onloadend = () => {
                const base64Data = (reader.result as string).split(',')[1];
                ws.send(JSON.stringify({
                    type: 'analyze',
                    file_data: base64Data,
                    file_type: file.type,
                    is_wip: true,
                    access_token: accessToken
                }));
            };
            reader.onerror = () => {
                ws.close();
                reject(new Error('Failed to read file'));
            };
            reader.readAsDataURL(file);
        };

        ws.onmessage = (event) => {
            try {
                const message = JSON.parse(event.data);

                if (message.type === 'agent_completed') {
                    onAgentUpdate(message.agent_name as AgentName, message.review);
                } else if (message.type === 'summary') {
                    onAgentUpdate('Summary');
                } else if (message.type === 'complete') {
                    resolved = true;
                    ws.close();
                    resolve(message.result?.leadAgentSummary || 'Analysis complete.');
                } else if (message.type === 'error') {
                    resolved = true;
                    ws.close();
                    reject(new Error(message.message || 'Analysis failed'));
                }
            } catch (e) {
                console.error('Failed to parse WebSocket message:', e);
            }
        };

        ws.onerror = () => {
            if (!resolved) {
                resolved = true;
                reject(new Error('WebSocket connection error'));
            }
        };

        ws.onclose = (event) => {
            if (!resolved && !event.wasClean) {
                resolved = true;
                reject(new Error('Connection closed unexpectedly'));
            }
        };
    });
};

/**
 * Get a chat response from the WIP Lead Agent.
 * Uses HTTP REST endpoint (not WebSocket).
 */
export const getWIPChatResponse = async (prompt: string): Promise<string> => {
    try {
        const response = await fetch(`${HTTP_URL}/api/chat`, {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify({ prompt })
        });

        if (!response.ok) {
            throw new Error(`HTTP ${response.status}: ${response.statusText}`);
        }

        const data = await response.json();
        return data.response || data.message || 'No response from Lead Agent.';
    } catch (error) {
        console.error('Error getting WIP chat response:', error);
        // Fallback message when backend chat endpoint is not available
        return "I'm sorry, the chat feature requires the backend chat endpoint to be implemented. For now, please use the proof analysis feature.";
    }
};