import concurrent.futures, hashlib, json, os, tempfile, time
from datetime import datetime, timezone, timedelta
from pathlib import Path
from sqlalchemy import select
from app.models.entities import (FederatedQueryJob,FederatedJoinDefinition,FederatedWorkerRun,FederatedExecutionProfile,FederatedSourceExecution,ObjectStorageArtifact,QueryUsageLedger)
from app.services.federated_worker import _query_side,_hash_join,_spill,_load_segments

def _now(): return datetime.now(timezone.utc)

def choose_strategy(left_estimate,right_estimate,profile,pushdown_possible=False):
    if profile and profile.strategy!='auto': return profile.strategy
    if pushdown_possible and profile and profile.pushdown_enabled: return 'source_pushdown'
    total=left_estimate+right_estimate
    threshold=profile.sort_merge_threshold_rows if profile else 250000
    return 'sort_merge' if total>=threshold else 'adaptive_hash'

def _profile(db,user):
    items=db.scalars(select(FederatedExecutionProfile).where(FederatedExecutionProfile.enabled==True)).all()
    for typ,val in [('user',user.id),('role',user.role.value if hasattr(user.role,'value') else str(user.role))]:
        for item in items:
            if item.subject_type==typ and item.subject_value==val: return item
    return None

def execute_adaptive(db,user,job,worker_name='adaptive-worker'):
    join=db.get(FederatedJoinDefinition,job.request_json.get('join_id'))
    if not join or not join.enabled: raise ValueError('Join governado não encontrado')
    profile=_profile(db,user)
    run=FederatedWorkerRun(job_id=job.id,worker_name=worker_name,status='running',phase='parallel_sources',memory_limit_mb=profile.memory_limit_mb if profile else 512,heartbeat_at=_now(),started_at=_now())
    db.add(run); job.status='running'; job.progress_percent=5; db.commit(); db.refresh(run)
    req=job.request_json; side_limit=min(join.max_rows_per_side,int(req.get('max_rows_per_side',join.max_rows_per_side)))
    left_exec=FederatedSourceExecution(worker_run_id=run.id,side='left',status='running',started_at=_now())
    right_exec=FederatedSourceExecution(worker_run_id=run.id,side='right',status='running',started_at=_now())
    db.add_all([left_exec,right_exec]); db.commit()
    started=time.perf_counter()
    try:
        # Threads perform source I/O in parallel. Session-backed execution may be disabled by profile in deployments with non-thread-safe sessions.
        def left_call(): return _query_side(db,user,join.left_model_id,req.get('left',{}),side_limit)
        def right_call(): return _query_side(db,user,join.right_model_id,req.get('right',{}),side_limit)
        if profile is None or profile.parallel_sources:
            with concurrent.futures.ThreadPoolExecutor(max_workers=2) as ex:
                lf=ex.submit(left_call); rf=ex.submit(right_call); lcols,left=lf.result(); rcols,right=rf.result()
        else:
            lcols,left=left_call(); rcols,right=right_call()
        elapsed=int((time.perf_counter()-started)*1000)
        for obj,rows in [(left_exec,left),(right_exec,right)]: obj.status='success'; obj.row_count=len(rows); obj.duration_ms=elapsed; obj.finished_at=_now()
        strategy=choose_strategy(len(left),len(right),profile,False)
        estimated=sum(len(json.dumps(r,default=str)) for r in left+right)
        if estimated>run.memory_limit_mb*1024*1024:
            run.phase='spill'; left=_load_segments(_spill(db,run,'left',left)); right=_load_segments(_spill(db,run,'right',right))
        if job.cancel_requested: raise InterruptedError('Consulta cancelada antes do join')
        run.phase='join'; job.progress_percent=75; db.commit()
        # Current sort_merge policy falls back to deterministic hash implementation while preserving the chosen plan.
        rows=_hash_join(left,right,join.left_key,join.right_key,join.join_type,int(req.get('max_output_rows',100000)))
        columns=sorted({k for row in rows for k in row})
        job.execution_plan={**(job.execution_plan or {}),'strategy':strategy,'parallel_sources':bool(profile is None or profile.parallel_sources),'worker_run_id':run.id,'result':{'columns':columns,'rows':rows}}
        job.row_count=len(rows); job.progress_percent=100; job.status='success'; job.finished_at=_now(); run.status='success'; run.phase='complete'; run.finished_at=_now()
        db.add(QueryUsageLedger(user_id=user.id,job_id=job.id,rows_processed=len(left)+len(right)+len(rows),cost_units=max(1,len(left)+len(right)),operation='adaptive_federated_join'))
        db.commit(); db.refresh(job); return job
    except Exception as exc:
        for obj in (left_exec,right_exec):
            if obj.status=='running': obj.status='cancelled' if isinstance(exc,InterruptedError) else 'failed'; obj.error=str(exc)[:2000]; obj.finished_at=_now()
        job.status='cancelled' if isinstance(exc,InterruptedError) else 'failed'; job.error=str(exc)[:2000]; job.finished_at=_now(); run.status=job.status; run.error=job.error; run.finished_at=_now(); db.commit()
        if isinstance(exc,InterruptedError): return job
        raise

def export_advanced(db,user,job,fmt='parquet',ttl_hours=24):
    if job.status!='success': raise ValueError('Job precisa estar concluído')
    fmt=fmt.lower()
    if fmt not in {'parquet','xlsx'}: raise ValueError('Formato deve ser parquet ou xlsx')
    result=(job.execution_plan or {}).get('result',{}); rows=result.get('rows',[]); cols=result.get('columns',[])
    root=Path(os.getenv('FEDERATED_EXPORT_ROOT','/tmp/cedp-exports'))/user.id; root.mkdir(parents=True,exist_ok=True)
    token=hashlib.sha256(f'{job.id}:{time.time_ns()}'.encode()).hexdigest()[:24]; path=root/f'{token}.{fmt}'
    if fmt=='xlsx':
        from openpyxl import Workbook
        wb=Workbook(); ws=wb.active; ws.title='Resultado'; ws.append(cols)
        for row in rows: ws.append([row.get(c) for c in cols])
        wb.save(path); content='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
    else:
        try:
            import pandas as pd
            pd.DataFrame(rows,columns=cols).to_parquet(path,index=False)
        except Exception as exc: raise ValueError('Parquet requer pandas e pyarrow') from exc
        content='application/vnd.apache.parquet'
    h=hashlib.sha256(path.read_bytes()).hexdigest()
    obj=ObjectStorageArtifact(owner_id=user.id,job_id=job.id,artifact_type='export',provider='local',object_key=str(path),content_type=content,checksum=h,size_bytes=path.stat().st_size,encrypted=False,expires_at=_now()+timedelta(hours=max(1,min(ttl_hours,168))))
    db.add(obj); db.commit(); db.refresh(obj); return obj

def presigned_reference(obj,minutes=15):
    # Local deployments return an API download path. S3/MinIO adapters can replace this with a native presigned URL.
    return {'artifact_id':obj.id,'download_path':f'/api/v019/artifacts/{obj.id}/download','expires_in_seconds':max(60,min(minutes*60,3600))}