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, AuditRepository
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,
    current_user_id: Optional[uuid.UUID] = None,
) -> 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

        # Compute file hash for duplicate detection
        file_hash = storage_service.get_checksum(file_data)
        logger.info(f"[WEBSOCKET] Computed file hash: {file_hash}")

        # 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,
                        "resolvedIssues": review.resolvedIssues,
                        "outstandingIssues": review.outstandingIssues,
                        "newIssues": review.newIssues,
                    }
                })

        # Extract brand selection for brand agent
        brand = data.get("brand", "Barclaycard")  # Default to Barclaycard if not specified
        logger.info(f"[WEBSOCKET] Brand selection: {brand}")

        # Fetch previous analysis if this is a revision
        previous_analysis = None
        previous_file_hash = None
        is_identical_file = False
        campaign_id = data.get("campaign_id")
        proof_name = data.get("proof_name")

        if campaign_id and proof_name:
            try:
                logger.info(f"[WEBSOCKET] Checking for previous analysis - campaign: {campaign_id}, proof: {proof_name}")
                async with async_session_factory() as session:
                    proof_repo = ProofRepository(session)
                    existing_proof = await proof_repo.get_by_campaign_and_name(
                        uuid.UUID(campaign_id), proof_name
                    )
                    if existing_proof:
                        previous_analysis = await proof_repo.get_latest_version_review(existing_proof.id)
                        previous_file_hash = await proof_repo.get_latest_version_hash(existing_proof.id)
                        if previous_analysis:
                            logger.info(f"[WEBSOCKET] Found previous analysis version {previous_analysis.get('version')}")
                        else:
                            logger.info("[WEBSOCKET] No previous analysis found (new proof)")
                        # Check if file is identical to previous version
                        if previous_file_hash and previous_file_hash == file_hash:
                            is_identical_file = True
                            logger.info(f"[WEBSOCKET] Identical file detected - hash matches previous version: {file_hash}")
                    else:
                        logger.info("[WEBSOCKET] No existing proof found (new proof)")
            except Exception as e:
                logger.warning(f"[WEBSOCKET] Failed to fetch previous analysis: {str(e)}")
                # Continue without previous analysis - still run the current analysis

        # Extract proof metadata for agent context
        channel = data.get("channel")
        sub_channel = data.get("sub_channel")
        proof_type = data.get("proof_type")

        # Build a once-only callback that notifies the client when the primary
        # Gemini model is unavailable and the fallback model is used instead.
        fallback_notified = False

        async def on_model_fallback() -> None:
            nonlocal fallback_notified
            if fallback_notified:
                return
            fallback_notified = True
            logger.info(f"[WEBSOCKET] Sending model_fallback notification to client: {client_id}")
            if manager.is_connected(client_id):
                await manager.send_message(client_id, {"type": "model_fallback"})
            else:
                logger.warning(f"[WEBSOCKET] Cannot send model_fallback - client {client_id} not connected")

        # Run the analysis
        logger.info("[WEBSOCKET] Starting analysis...")
        result, pdf_pages = await analysis_service.analyze_proof(
            file_data=file_data,
            file_type=file_type,
            on_agent_update=on_agent_update,
            is_wip=is_wip,
            brand=brand,
            previous_analysis=previous_analysis,
            channel=channel,
            sub_channel=sub_channel,
            proof_type=proof_type,
            on_fallback=on_model_fallback,
        )

        # 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,
                "resolvedIssues": result.legalAgentReview.resolvedIssues,
                "outstandingIssues": result.legalAgentReview.outstandingIssues,
                "newIssues": result.legalAgentReview.newIssues,
            },
            "brandAgentReview": {
                "ragStatus": result.brandAgentReview.ragStatus,
                "feedback": result.brandAgentReview.feedback,
                "issues": result.brandAgentReview.issues,
                "resolvedIssues": result.brandAgentReview.resolvedIssues,
                "outstandingIssues": result.brandAgentReview.outstandingIssues,
                "newIssues": result.brandAgentReview.newIssues,
            },
            "channelBestPracticesAgentReview": {
                "ragStatus": result.channelBestPracticesAgentReview.ragStatus,
                "feedback": result.channelBestPracticesAgentReview.feedback,
                "issues": result.channelBestPracticesAgentReview.issues,
                "resolvedIssues": result.channelBestPracticesAgentReview.resolvedIssues,
                "outstandingIssues": result.channelBestPracticesAgentReview.outstandingIssues,
                "newIssues": result.channelBestPracticesAgentReview.newIssues,
            },
            "channelTechSpecsAgentReview": {
                "ragStatus": result.channelTechSpecsAgentReview.ragStatus,
                "feedback": result.channelTechSpecsAgentReview.feedback,
                "issues": result.channelTechSpecsAgentReview.issues,
                "resolvedIssues": result.channelTechSpecsAgentReview.resolvedIssues,
                "outstandingIssues": result.channelTechSpecsAgentReview.outstandingIssues,
                "newIssues": result.channelTechSpecsAgentReview.newIssues,
            },
            "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

        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 images up to 10MB (data URLs are ~33% larger than binary)
                    thumbnail_url = None
                    if len(file_data) < 10000000:
                        if file_type.startswith('image/'):
                            thumbnail_url = await storage_service.generate_thumbnail_data_url(file_data, file_type)
                        elif file_type == 'application/pdf' and pdf_pages:
                            # Use first rasterized page as thumbnail
                            first_page_data, _, _ = pdf_pages[0]
                            thumbnail_url = await storage_service.generate_thumbnail_data_url(first_page_data, 'image/png')

                    # 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,
                        file_hash=file_hash,
                        is_identical_file=is_identical_file,
                        created_by=current_user_id,
                    )

                    # Auto-create ErrorItem when analysis results in an error
                    if result.overallStatus == "Analysis Error":
                        audit_repo = AuditRepository(session)
                        await audit_repo.create_error_item(
                            proof_version_id=version.id,
                            error_summary=result.leadAgentSummary,
                        )
                        logger.info(f"[WEBSOCKET] Created ErrorItem for analysis error on version {version.id}")

                    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,
                "is_identical_file": is_identical_file,
            }
            # Include proof/version IDs if persisted
            if proof_id:
                response["proof_id"] = proof_id
            if version_id:
                response["version_id"] = version_id

            # Include rasterized PDF pages if present
            if pdf_pages:
                import base64 as b64_module
                response["pdf_pages"] = [
                    {
                        "page": i + 1,
                        "data_url": f"data:image/png;base64,{b64_module.b64encode(png_data).decode('utf-8')}",
                        "width": width,
                        "height": height,
                    }
                    for i, (png_data, width, height) in enumerate(pdf_pages)
                ]
                logger.info(f"[WEBSOCKET] Including {len(pdf_pages)} rasterized PDF pages in response")

            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)}"
            })
