import base64
import logging
import uuid
from typing import Optional

from fastapi import WebSocket

from app.websocket.manager import ConnectionManager
from app.services.analysis_service import AnalysisService
from app.models.schemas import SubReview
from app.models.database import async_session_factory
from app.repositories import ProofRepository, CampaignRepository, UserRepository
from app.services.storage_service import storage_service

logger = logging.getLogger(__name__)


async def handle_analyze_message(
    websocket: WebSocket,
    client_id: str,
    data: dict,
    manager: ConnectionManager,
    analysis_service: AnalysisService,
) -> None:
    """
    Handle an 'analyze' message from the client.

    Runs the proof analysis and sends real-time updates via WebSocket.

    Args:
        websocket: The WebSocket connection
        client_id: Unique client identifier
        data: The message data containing file_data, file_type, is_wip
        manager: Connection manager for sending messages
        analysis_service: Service to run the analysis
    """
    try:
        logger.info(f"[WEBSOCKET] Received analyze request from client: {client_id}")

        # Extract and decode the file data
        file_data_b64 = data.get("file_data", "")
        file_type = data.get("file_type", "image/png")
        is_wip = data.get("is_wip", False)

        logger.info(f"[WEBSOCKET] File type: {file_type}, is_wip: {is_wip}, base64 length: {len(file_data_b64)}")

        # Decode base64 file data
        try:
            file_data = base64.b64decode(file_data_b64)
            logger.info(f"[WEBSOCKET] Decoded file size: {len(file_data)} bytes")
        except Exception as e:
            logger.error(f"[WEBSOCKET] Failed to decode file data: {str(e)}")
            await manager.send_message(client_id, {
                "type": "error",
                "message": f"Failed to decode file data: {str(e)}"
            })
            return

        # Create callback for real-time updates
        async def on_agent_update(agent_name: str, review: SubReview | None) -> None:
            if not manager.is_connected(client_id):
                logger.warning(f"[WEBSOCKET] Client {client_id} disconnected, skipping update")
                return

            if review is None:
                # Agent is starting
                logger.info(f"[WEBSOCKET] Sending agent_started: {agent_name}")
                await manager.send_message(client_id, {
                    "type": "agent_started",
                    "agent_name": agent_name
                })
            else:
                # Agent completed
                logger.info(f"[WEBSOCKET] Sending agent_completed: {agent_name} - ragStatus: {review.ragStatus}")
                await manager.send_message(client_id, {
                    "type": "agent_completed",
                    "agent_name": agent_name,
                    "review": {
                        "ragStatus": review.ragStatus,
                        "feedback": review.feedback,
                        "issues": review.issues,
                        "isFinancialPromotion": review.isFinancialPromotion,
                        "financialPromotionReason": review.financialPromotionReason,
                    }
                })

        # Run the analysis
        logger.info("[WEBSOCKET] Starting analysis...")
        result = await analysis_service.analyze_proof(
            file_data=file_data,
            file_type=file_type,
            on_agent_update=on_agent_update,
            is_wip=is_wip,
        )

        # Build the result dict
        result_dict = {
            "legalAgentReview": {
                "ragStatus": result.legalAgentReview.ragStatus,
                "feedback": result.legalAgentReview.feedback,
                "issues": result.legalAgentReview.issues,
                "isFinancialPromotion": result.legalAgentReview.isFinancialPromotion,
                "financialPromotionReason": result.legalAgentReview.financialPromotionReason,
            },
            "brandAgentReview": {
                "ragStatus": result.brandAgentReview.ragStatus,
                "feedback": result.brandAgentReview.feedback,
                "issues": result.brandAgentReview.issues,
            },
            "toneAgentReview": {
                "ragStatus": result.toneAgentReview.ragStatus,
                "feedback": result.toneAgentReview.feedback,
                "issues": result.toneAgentReview.issues,
            },
            "channelAgentReview": {
                "ragStatus": result.channelAgentReview.ragStatus,
                "feedback": result.channelAgentReview.feedback,
                "issues": result.channelAgentReview.issues,
            },
            "leadAgentSummary": result.leadAgentSummary,
            "overallStatus": result.overallStatus,
            "financialPromotionReason": result.financialPromotionReason,
        }

        # Persist to database if campaign info provided
        proof_id: Optional[str] = None
        version_id: Optional[str] = None
        campaign_id = data.get("campaign_id")
        proof_name = data.get("proof_name")

        if campaign_id and proof_name:
            try:
                logger.info(f"[WEBSOCKET] Persisting result for campaign {campaign_id}, proof {proof_name}")
                async with async_session_factory() as session:
                    proof_repo = ProofRepository(session)

                    # Store the file
                    file_storage_key = await storage_service.store_file(
                        file_data=file_data,
                        campaign_id=uuid.UUID(campaign_id),
                        proof_name=proof_name,
                        version=1,  # Will be updated by add_version_with_review
                        file_type=file_type,
                    )

                    # Generate thumbnail for small files
                    thumbnail_url = None
                    if len(file_data) < 500000:  # < 500KB
                        thumbnail_url = await storage_service.generate_thumbnail_data_url(file_data, file_type)

                    # Save proof and version
                    proof, version = await proof_repo.add_version_with_review(
                        campaign_id=uuid.UUID(campaign_id),
                        proof_name=proof_name,
                        channel=data.get("channel"),
                        sub_channel=data.get("sub_channel"),
                        proof_type=data.get("proof_type"),
                        file_storage_key=file_storage_key,
                        thumbnail_url=thumbnail_url,
                        agent_review=result_dict,
                        overall_status=result.overallStatus,
                    )

                    await session.commit()
                    proof_id = str(proof.id)
                    version_id = str(version.id)
                    logger.info(f"[WEBSOCKET] Persisted proof {proof_id} version {version.version}")

            except Exception as e:
                logger.error(f"[WEBSOCKET] Failed to persist result: {str(e)}")
                # Continue - still send result to client even if persistence fails

        # Send the complete result
        logger.info(f"[WEBSOCKET] Analysis complete, sending result - overallStatus: {result.overallStatus}")
        if manager.is_connected(client_id):
            response = {
                "type": "complete",
                "result": result_dict,
            }
            # Include proof/version IDs if persisted
            if proof_id:
                response["proof_id"] = proof_id
            if version_id:
                response["version_id"] = version_id

            await manager.send_message(client_id, response)
            logger.info(f"[WEBSOCKET] Result sent to client: {client_id}")

    except Exception as e:
        # Send error message
        logger.error(f"[WEBSOCKET] Analysis failed for client {client_id}: {str(e)}")
        if manager.is_connected(client_id):
            await manager.send_message(client_id, {
                "type": "error",
                "message": f"Analysis failed: {str(e)}"
            })
