app/tasks.py

import asyncio
import json
from datetime import datetime, timedelta
import redis
from sqlalchemy import update
from sqlalchemy.future import select
from app.config import settings
from contextlib import asynccontextmanager

_task_engine = None
_task_session_factory = None

def get_task_session_factory():
    global _task_engine, _task_session_factory
    if _task_session_factory is None:
        from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
        from uuid import uuid4
        database_url = settings.DATABASE_URL
        if database_url.startswith("postgresql://"):
            database_url = database_url.replace("postgresql://", "postgresql+asyncpg://", 1)
        _task_engine = create_async_engine(
            database_url,
            pool_size=5,
            max_overflow=10,
            pool_pre_ping=True,
            connect_args={
                "statement_cache_size": 0,
                "prepared_statement_cache_size": 0,
                "prepared_statement_name_func": lambda: f"__asyncpg_{uuid4().hex}__"
            }
        )
        _task_session_factory = async_sessionmaker(bind=_task_engine, class_=AsyncSession, expire_on_commit=False)
    return _task_session_factory

@asynccontextmanager
async def local_session():
    factory = get_task_session_factory()
    async with factory() as session:
        yield session
from app import models
from app.utils.ai import run_code_audit, generate_system_design
from loguru import logger

# Redis client for publishing streaming progress events
redis_client = redis.Redis.from_url(settings.REDIS_URL)

def publish_event(job_id: str, event_type: str, data: dict = None):
    payload = {"event": event_type}
    if data:
        payload.update(data)
    raw_payload = json.dumps(payload)
    redis_client.publish(f"job:{job_id}", raw_payload)
    try:
        redis_client.rpush(f"job:{job_id}:events", raw_payload)
        redis_client.expire(f"job:{job_id}:events", 3600)
    except Exception as e:
        logger.error(f"Failed to cache event in Redis: {e}")

async def run_audit_task(project_id_str: str, job_id: str, files: list, file_type: str):
    logger.info(f"Starting audit task for project {project_id_str}, job {job_id}")
    
    # Publish initial agent start events
    publish_event(job_id, "agent_started", {"agent": "connecting"})
    publish_event(job_id, "agent_complete", {"agent": "connecting"})
    publish_event(job_id, "agent_started", {"agent": "sre"})
    publish_event(job_id, "agent_started", {"agent": "backend"})
    publish_event(job_id, "agent_started", {"agent": "infrastructure"})
    publish_event(job_id, "agent_started", {"agent": "cloud_architect"})
    
    try:
        # Run AI audit analysis
        report_data = await run_code_audit(files, file_type)
        
        # Calculate stats for confidence score
        files_count = len(files)
        total_chars = sum(len(f["content"]) for f in files)
        
        has_config_files = any(
            any(f["filename"].endswith(ext) for ext in [".json", ".yaml", ".yml", ".toml", ".config", ".env"])
            for f in files
        )
        has_docker = any("Dockerfile" in f["filename"] for f in files)
        
        # Confidence score weighting:
        # files>=10(+20), config_files(+15), docker(+15), chars>=50K(+20), not_truncated(+20)
        confidence_score = 0
        if files_count >= 10:
            confidence_score += 20
        if has_config_files:
            confidence_score += 15
        if has_docker:
            confidence_score += 15
        if total_chars >= 50000:
            confidence_score += 20
        confidence_score += 20  # Assuming not truncated by default
        
        if confidence_score >= 70:
            confidence_level = "high"
        elif confidence_score >= 40:
            confidence_level = "medium"
        else:
            confidence_level = "low"
            
        confidence_dict = {
            "level": confidence_level,
            "score": confidence_score,
            "label": f"{confidence_level.capitalize()} Confidence",
            "based_on": report_data.get("confidence", {}).get("based_on", ["Initial file scan completed"]),
            "limitations": report_data.get("confidence", {}).get("limitations", ["No runtime metrics available"]),
            "to_increase_confidence": report_data.get("confidence", {}).get("to_increase_confidence", ["Upload load test scripts"])
        }
        
        # Enforce server-side score calculation
        # overall_score = round(sre*0.30 + backend*0.30 + infra*0.20 + cloud*0.20)
        sre_score = report_data.get("sre_score", 80)
        backend_score = report_data.get("backend_score", 80)
        infra_score = report_data.get("infra_score", 80)
        cloud_score = report_data.get("cloud_score", 80)
        
        overall_score = round(
            sre_score * 0.30 +
            backend_score * 0.30 +
            infra_score * 0.20 +
            cloud_score * 0.20
        )
        
        # Enforce capacity tiers
        # 0-30→50-200 · 31-50→200-1K · 51-65→1K-5K · 66-79→5K-25K · 80-89→25K-100K · 90-100→100K-500K
        if overall_score <= 30:
            safe, peak = "50-200 DAU", "200-1K DAU"
        elif overall_score <= 50:
            safe, peak = "200-1K DAU", "1K-5K DAU"
        elif overall_score <= 65:
            safe, peak = "1K-5K DAU", "5K-25K DAU"
        elif overall_score <= 79:
            safe, peak = "5K-25K DAU", "25K-100K DAU"
        elif overall_score <= 89:
            safe, peak = "25K-100K DAU", "100K-500K DAU"
        else:
            safe, peak = "100K-500K DAU", "500K-2M DAU"
            
        capacity_dict = {
            "safe_range": safe,
            "peak_range": peak,
            "description": report_data.get("capacity_estimate", {}).get("description", "DAU estimate based on static architecture scanning."),
            "reasoning": report_data.get("capacity_estimate", {}).get("reasoning", "Database concurrency constraints set maximum peaks."),
            "confidence": report_data.get("capacity_estimate", {}).get("confidence", "Medium")
        }
        
        # Prepare Agent reports mapping
        agents_data = report_data.get("agents", {})
        
        # Re-set agents overall scores if modified
        if "sre" in agents_data:
            agents_data["sre"]["score"] = sre_score
        if "backend" in agents_data:
            agents_data["backend"]["score"] = backend_score
        if "infrastructure" in agents_data:
            agents_data["infrastructure"]["score"] = infra_score
        if "cloud_architect" in agents_data:
            agents_data["cloud_architect"]["score"] = cloud_score
            
        disclaimer = (
            "This score is based on static code analysis only. It does "
            "not guarantee runtime performance. Validate with load testing "
            "before production launch."
        )
        
        async with local_session() as session:
            # Load project
            import uuid
            project_uuid = uuid.UUID(project_id_str)
            proj_result = await session.execute(select(models.Project).where(models.Project.id == project_uuid))
            project = proj_result.scalars().first()
            
            if not project:
                raise Exception("Project not found")
                
            # Create Audit Report
            audit_report = models.AuditReport(
                project_id=project.id,
                overall_score=overall_score,
                confidence=confidence_dict,
                capacity_estimate=capacity_dict,
                agents=agents_data,
                top_critical_issues=report_data.get("top_critical_issues", []),
                quick_wins=report_data.get("quick_wins", []),
                benchmark_percentile=report_data.get("benchmark_percentile", 75.0),
                score_disclaimer=disclaimer,
                files_analyzed=files_count,
                files_skipped=0,
                was_truncated=False
            )
            
            session.add(audit_report)
            project.status = "complete"
            
            # Update user stats
            user_result = await session.execute(select(models.User).where(models.User.id == project.user_id))
            user = user_result.scalars().first()
            if user:
                user.total_audits += 1
                session.add(user)
                
            await session.commit()
            
        # Complete agents progress steps
        publish_event(job_id, "agent_complete", {"agent": "sre"})
        publish_event(job_id, "agent_complete", {"agent": "backend"})
        publish_event(job_id, "agent_complete", {"agent": "infrastructure"})
        publish_event(job_id, "agent_complete", {"agent": "cloud_architect"})
        
        publish_event(job_id, "job_complete", {"audit_id": str(project.id)})
        logger.info(f"Audit job {job_id} completed successfully")
        
    except Exception as e:
        logger.error(f"Audit task failed: {e}")
        async with local_session() as session:
            import uuid
            project_uuid = uuid.UUID(project_id_str)
            proj_result = await session.execute(select(models.Project).where(models.Project.id == project_uuid))
            project = proj_result.scalars().first()
            if project:
                project.status = "failed"
                project.error_msg = str(e)
                await session.commit()
        publish_event(job_id, "job_failed", {"error": str(e)})

async def run_design_task(project_id_str: str, job_id: str, idea_prompt: str):
    logger.info(f"Starting design task for project {project_id_str}, job {job_id}")
    
    # Emit progress events without artificial delay
    publish_event(job_id, "agent_started", {"agent": "connecting"})
    publish_event(job_id, "agent_complete", {"agent": "connecting"})
    publish_event(job_id, "agent_started", {"agent": "sre"})
    publish_event(job_id, "agent_started", {"agent": "backend"})
    publish_event(job_id, "agent_started", {"agent": "infrastructure"})
    publish_event(job_id, "agent_started", {"agent": "cloud_architect"})
    
    try:
        # Run AI System Design
        design_data = await generate_system_design(idea_prompt)
        
        async with local_session() as session:
            # Load project
            import uuid
            project_uuid = uuid.UUID(project_id_str)
            proj_result = await session.execute(select(models.Project).where(models.Project.id == project_uuid))
            project = proj_result.scalars().first()
            
            if not project:
                raise Exception("Project not found")
                
            # Create System Design
            system_design = models.SystemDesign(
                project_id=project.id,
                idea_prompt=idea_prompt,
                title=design_data.get("title", "System Architecture Blueprint"),
                founder_summary=design_data.get("founder_summary", ""),
                engineer_summary=design_data.get("engineer_summary", ""),
                architecture_type=design_data.get("architecture_type", "Hybrid"),
                reasoning=design_data.get("reasoning", ""),
                stack=design_data.get("stack", {}),
                database_design=design_data.get("database_design", {}),
                api_design=design_data.get("api_design", {}),
                infrastructure=design_data.get("infrastructure", {}),
                reliability=design_data.get("reliability", {}),
                cost_estimates=design_data.get("cost_estimates", {}),
                diagram=design_data.get("diagram", {"nodes": [], "edges": []})
            )
            
            session.add(system_design)
            project.status = "complete"
            
            # Update user stats
            user_result = await session.execute(select(models.User).where(models.User.id == project.user_id))
            user = user_result.scalars().first()
            if user:
                user.total_designs += 1
                session.add(user)
                
            await session.commit()
            
        publish_event(job_id, "agent_complete", {"agent": "cloud_architect"})
        publish_event(job_id, "job_complete", {"design_id": str(project.id)})
        logger.info(f"Design job {job_id} completed successfully")
        
    except Exception as e:
        logger.error(f"Design task failed: {e}")
        async with local_session() as session:
            import uuid
            project_uuid = uuid.UUID(project_id_str)
            proj_result = await session.execute(select(models.Project).where(models.Project.id == project_uuid))
            project = proj_result.scalars().first()
            if project:
                project.status = "failed"
                project.error_msg = str(e)
                await session.commit()
        publish_event(job_id, "job_failed", {"error": str(e)})

# RQ requires functions to be importable at module level (no inner/local functions)
# This is the sync wrapper for GitHub audit jobs — called by RQ worker
def run_github_audit_job(proj_id: str, job_id: str, repo: str, branch: str, access_token: str, file_type: str):
    """Module-level RQ job for GitHub-sourced audits. Picklable by RQ."""
    from app.utils.github_client import get_repo_files
    files = get_repo_files(access_token, repo, branch)
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        loop.run_until_complete(run_audit_task(proj_id, job_id, files, file_type))
    finally:
        loop.close()

# Sync wrappers for RQ — each creates its own event loop to avoid cross-thread conflicts
def audit_job_wrapper(project_id_str: str, job_id: str, files: list, file_type: str):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        loop.run_until_complete(run_audit_task(project_id_str, job_id, files, file_type))
    finally:
        loop.close()

def design_job_wrapper(project_id_str: str, job_id: str, idea_prompt: str):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        loop.run_until_complete(run_design_task(project_id_str, job_id, idea_prompt))
    finally:
        loop.close()