
    iBX                       d Z ddlmZ ddlZddlZddlZddlZddlmZm	Z	 ddl
mZ ddlmZmZmZmZmZmZ ddlmZ  ej        e          ZdZd	Zd
ZdZdZdZdZdZdZ dZ!dZ"dZ# G d de$          Z% G d de$          Z&e G d d                      Z'e G d d                      Z(e G d d                      Z)d5d Z*ed!eee+ef                  f         Z,	  G d" d#          Z-d6d'Z.d7d,Z/d8d/Z0d9d4Z1dS ):u  QQ Bot chunked upload flow.

The QQ v2 API caps inline base64 uploads (``file_data`` / ``url``) at ~10 MB.
For files between 10 MB and ~100 MB we have to use the three-step chunked
upload flow::

    1. POST /v2/{users|groups}/{id}/upload_prepare
       → returns upload_id, block_size, and an array of pre-signed COS part URLs.
    2. For each part:
         PUT the part bytes to its pre-signed COS URL,
         then POST /v2/{users|groups}/{id}/upload_part_finish to acknowledge.
    3. POST /v2/{users|groups}/{id}/files with {"upload_id": ...}
       → returns the ``file_info`` token the caller uses in a RichMedia
       message.

Error-code semantics (from the QQ Bot v2 API spec):

- ``40093001`` — ``upload_part_finish`` retryable. Retry until the server-provided
  ``retry_timeout`` elapses (or a local cap).
- ``40093002`` — daily cumulative upload quota exceeded. Not retryable; surface
  as :class:`UploadDailyLimitExceededError` so the caller can build a
  user-friendly reply.

Exceptions:

- :class:`UploadDailyLimitExceededError` — daily quota hit (non-retryable).
- :class:`UploadFileTooLargeError` — file exceeds the platform per-file limit.
- :class:`RuntimeError` — generic upload failure (network, part PUT, complete).

Ported from WideLee's qqbot-agent-sdk v1.2.2 (``media_loader.py::ChunkedUploader``)
so the heavy-upload path stays in-tree. Authorship preserved via Co-authored-by.
    )annotationsN)	dataclassfield)Path)Any	AwaitableCallableDictListOptional)FILE_UPLOAD_TIMEOUTiJciIc   
   g     r@         ?g      ^@g     @g       @i  c                  >     e Zd ZdZdd fd
Zedd            Z xZS )UploadDailyLimitExceededErrorzRaised when ``upload_prepare`` returns biz_code 40093002.

    The daily cumulative upload quota for this bot has been reached. Callers
    should surface :attr:`file_name` + :attr:`file_size_human` so the model
    can compose a helpful reply.
     	file_namestr	file_sizeintmessagereturnNonec                p    || _         || _        t                                          |pd|           d S )Nz Daily upload limit exceeded for )r   r   super__init__)selfr   r   r   	__class__s       K/home/piyush/.hermes/hermes-agent/gateway/platforms/qqbot/chunked_upload.pyr   z&UploadDailyLimitExceededError.__init__P   sH    ""GG)GG	
 	
 	
 	
 	
    c                *    t          | j                  S Nformat_sizer   r   s    r!   file_size_humanz-UploadDailyLimitExceededError.file_size_humanW       4>***r"   )r   )r   r   r   r   r   r   r   r   r   r   )__name__
__module____qualname____doc__r   propertyr(   __classcell__r    s   @r!   r   r   H   sl         
 
 
 
 
 
 
 + + + X+ + + + +r"   r   c                  Z     e Zd ZdZ	 	 dd fdZedd            Zedd            Z xZS )UploadFileTooLargeErrorz<Raised when a file exceeds the platform per-file size limit.r   r   r   r   r   r   limit_bytesr   r   r   c                    || _         || _        || _        |rdt          |           dnd}t	                                          |pd|dt          |           d|            d S )Nz ()r   zFile z) exceeds platform limit)r   r   r4   r&   r   r   )r   r   r   r4   r   	limit_strr    s         r!   r   z UploadFileTooLargeError.__init___   s     #"&8CK4[114444	 5	 5 5{9'='= 5 5)25 5	
 	
 	
 	
 	
r"   c                *    t          | j                  S r$   r%   r'   s    r!   r(   z'UploadFileTooLargeError.file_size_humanr   r)   r"   c                <    | j         rt          | j                   ndS )Nunknown)r4   r&   r'   s    r!   limit_humanz#UploadFileTooLargeError.limit_humanv   s     040@O{4+,,,iOr"   )r   r   )
r   r   r   r   r4   r   r   r   r   r   r*   )	r+   r,   r-   r.   r   r/   r(   r;   r0   r1   s   @r!   r3   r3   \   s        FF 
 
 
 
 
 
 
& + + + X+ P P P XP P P P Pr"   r3   c                  H    e Zd ZU dZded<   dZded<   dZded<   dZded<   dS )_UploadProgressr   r   total_partstotal_bytescompleted_partsuploaded_bytesN)r+   r,   r-   r>   __annotations__r?   r@   rA    r"   r!   r=   r=   }   sV         KKONr"   r=   c                  2    e Zd ZU ded<   ded<   dZded<   dS )_PreparePartr   indexr   presigned_urlr   
block_sizeN)r+   r,   r-   rB   rH   rC   r"   r!   rE   rE      s8         JJJJr"   rE   c                  J    e Zd ZU ded<   ded<   ded<   eZded<   dZd	ed
<   dS )_PrepareResultr   	upload_idr   rH   zList[_PreparePart]partsconcurrency        floatretry_timeoutN)r+   r,   r-   rB   _DEFAULT_CONCURRENT_PARTSrM   rP   rC   r"   r!   rJ   rJ      sR         NNNOOO0K0000Mr"   rJ   rawDict[str, Any]r   c                   t          |                     d          t                    r|                     d          n| }t          |                    dd                    }|s't	          dt          |           dd                    t          |                    dd                    }|                    d	          p|                    d
          pg }t          |t                    r|s't	          dt          |           dd                    g }|D ]}t          |t                    s|                    t          t          |                    d          p|                    d          pd          t          |                    d          p|                    d          pd          t          |                    dd                                         t          |||t          |                    dt                              pt          t          |                    dd          pd                    S )zParse the upload_prepare API response into a normalized shape.

    The API may return the response directly or wrapped in ``data``.
    datarK   r   z+upload_prepare response missing upload_id: N   rH   r   rL   	part_listz'upload_prepare response missing parts: 
part_indexrF   rG   url)rF   rG   rH   rM   rP   rN   )rK   rH   rL   rM   rP   )
isinstancegetdictr   
ValueErrorr   listappendrE   rJ   rQ   rO   )rR   srcrK   rH   	raw_partsrL   ps          r!   _parse_prepare_responserc      s4   
 (>>
G#''&///CCCGGK,,--I 
J#c((4C4.JJ
 
 	
 SWW\1--..J  >CGGK$8$8>BIi&& 
i 
Fc#hhttnFF
 
 	
 !#E 
 
!T"" 	!%%--DwD1EE!EE/**@aeeEll@b  quu\15566  	
 	
 	
 	
 /HIIJJgNgCGGOS99@SAA   r"   .c                  N    e Zd ZdZ	 d-d.dZd/dZd0dZd1d!Zd2d'Zd3d*Z	d4d+Z
d,S )5ChunkedUploaderu  Run the prepare → PUT parts → complete sequence.

    :param api_request: Bound ``_api_request(method, path, body=..., timeout=...)``
        coroutine from the adapter. Must raise ``RuntimeError`` with the biz_code
        embedded in the message on API errors.
    :param http_put: Coroutine ``(url, data, headers, timeout) -> response`` for
        COS part uploads. Typically wraps ``httpx.AsyncClient.put``.
    :param log_tag: Log prefix.
    QQBotapi_requestApiRequestFnhttp_putCallable[..., Awaitable[Any]]log_tagr   r   r   c                0    || _         || _        || _        d S r$   )_api_request	_http_put_log_tag)r   rg   ri   rk   s       r!   r   zChunkedUploader.__init__   s     (!r"   	chat_type	target_id	file_path	file_typer   r   rS   c           	        
K   dvrt          d          t                    }|                                j        
t                              d j        |t          
          |           t          j	                    
                    dt          
           d{V }                     ||
|           d{V t          j        t                    }t          j        dk    rj        nt"          t$                    t                              d j        j        t          j                  t+          j                  |           t/          t+          j                  
          
 fdj        D             }	t1          |	|           d{V  t                              d	 j        t+          j                                                  j                   d{V S )
u  Run the full chunked upload and return the ``complete_upload`` response.

        :param chat_type: ``'c2c'`` or ``'group'``.
        :param target_id: User or group openid.
        :param file_path: Absolute path to a local file.
        :param file_type: ``MEDIA_TYPE_*`` constant.
        :param file_name: Original filename (for upload_prepare).
        :returns: The raw response dict from ``complete_upload`` — contains
            ``file_info`` that the caller uses in a RichMedia message body.
        :raises UploadDailyLimitExceededError: On biz_code 40093002.
        :raises UploadFileTooLargeError: When the file exceeds the platform limit.
        :raises RuntimeError: On other API or I/O failures.
        )c2cgroupz'ChunkedUploader: unsupported chat_type z2[%s] Chunked upload start: file=%s size=%s type=%dNr   zA[%s] Prepared: upload_id=%s block_size=%s parts=%d concurrency=%d)r>   r?   c                j    g | ]/}t          j        j        	j        j        | 
  
        0S ))	rp   rq   rr   r   rK   rsp_block_sizepartrP   progress)	functoolspartial_upload_one_partrK   rH   )
.0ry   rp   rr   r   preparerz   rP   r   rq   s
     r!   
<listcomp>z*ChunkedUploader.upload.<locals>.<listcomp>  sd     6
 6
 6
  %####!+&1+!  6
 6
 6
r"   u)   [%s] All %d parts uploaded, completing…)r]   r   statst_sizeloggerinforo   r&   asyncioget_running_looprun_in_executor_compute_file_hashes_prepareminrM   _MAX_CONCURRENT_PARTSrP   _PART_FINISH_DEFAULT_TIMEOUT_PART_FINISH_MAX_TIMEOUTrK   rH   lenrL   r=   _run_with_concurrency	_complete)r   rp   rq   rr   rs   r   pathhashesmax_concurrenttasksr   r   rz   rP   s   ````      @@@@r!   uploadzChunkedUploader.upload   s{     * ,,,G)GG   IIIKK'	@M9k)&<&<i	
 	
 	
 /11AA&	9
 
 
 
 
 
 
 

 y)Y	6
 
 
 
 
 
 
 
 W02GHH%,%:Q%>%>G!!D`$
 
 	OM7,k':L.M.M	
 	
 	
 #GM**!
 
 
6
 6
 6
 6
 6
 6
 6
 6
 6
 6
 6
  6
 6
 6
 $E>:::::::::7M3w}--	
 	
 	
 ^^Iy':KLLLLLLLLLr"   r   r   Dict[str, str]rJ   c                F  K   |dk    rdnd}| d| d}||||d         |d         |d         d	}		 |                      d
||	t                     d {V }
n># t          $ r1}t          |          }t           |v rt          |||          | d }~ww xY wt          |
          S )Nru   	/v2/users
/v2/groups/z/upload_preparemd5sha1md5_10m)rs   r   r   r   r   r   POSTbodytimeout)rm   r   RuntimeErrorr   _BIZ_CODE_DAILY_LIMITr   rc   )r   rp   rq   rs   r   r   r   baser   r   rR   excerr_msgs                r!   r   zChunkedUploader._prepare6  s      (500{{l33333"""%=6Ni(
 

	))41D *        CC  	 	 	#hhG')W443y'  	 's+++s   $A 
B ,BBrK   rx   ry   rE   rP   rO   rz   r=   c
           	       K   |j         }
|j        dk    r|j        n|}|
dz
  |z  }t          |||z
            }t          j                                        dt          |||           d{V }t          j        |          	                                }t                              d| j        |
|	j        t          |          ||           |                     |j        ||
|	j                   d{V  |                     ||||
|||           d{V  |	xj        dz  c_        |	xj        |z  c_        t                              d| j        |
|	j        |	j        |	j                   dS )z6PUT one part to COS, then call ``upload_part_finish``.r   r   Nz0[%s] Part %d/%d: uploading %s (offset=%d md5=%s)z"[%s] Part %d/%d done (%d/%d total))rF   rH   r   r   r   r   _read_file_chunkhashlibr   	hexdigestr   debugro   r>   r&   _put_to_presigned_urlrG   _part_finish_with_retryr@   rA   )r   rp   rq   rr   r   rK   rx   ry   rP   rz   rX   actual_block_sizeoffsetlengthrU   md5_hexs                   r!   r}   z ChunkedUploader._upload_one_partZ  s      Z
/3/B/BDOOq.N2&	F(:;; -//??"Ivv
 
 
 
 
 
 
 
 +d##--//>M:x';	
 	
 	
 ((j(2F
 
 	
 	
 	
 	
 	
 	
 	
 **y)
 
 	
 	
 	
 	
 	
 	
 	

 	  A%  6)0M:x';$h&:	
 	
 	
 	
 	
r"   rY   rU   bytesrX   r>   c                  K   d}t          t          dz             D ]M}	 t          j        |                     ||dt          t          |                    i          t                     d{V }t          |dd          }d|cxk    rd	k     r)n n&t          
                    d
| j        |||            dS d}		 t          |dd          dd         }	n# t          $ r Y nw xY wt          d| d|	           # t          $ rb}
|
}|t          k     rJdd|z  z  }t                              d| j        |||dz   ||
           t          j        |           d{V  Y d}
~
Gd}
~
ww xY wt          d| d| dt          dz    d|           )z1PUT part data to a pre-signed COS URL with retry.Nr   zContent-Length)rU   headers)r   status_coder   rV   i,  z[%s] PUT part %d/%d: %d OKr   textzCOS PUT returned z: r   r   z9[%s] PUT part %d/%d attempt %d failed, retry in %.1fs: %szPart r   z upload failed after  attempts: )range_PART_UPLOAD_MAX_RETRIESr   wait_forrn   r   r   _PART_UPLOAD_TIMEOUTgetattrr   r   ro   	Exceptionr   warningsleep)r   rY   rU   rX   r>   last_excattemptrespstatusbody_previewr   delays               r!   r   z%ChunkedUploader._put_to_presigned_url  sU      )-59:: #	/ #	/G"/$-NN!!13s4yy>> B #  
 1         !}a88&&&&&3&&&&&LL4z;   FF!#*4#<#<TcT#BLL    D"@@@,@@    	/ 	/ 	/5551<0ENNSz;!UC  
 "-.........	/ CJ C C C C'!+C C8@C C
 
 	
sC   BC<:C<=CC<
C$!C<#C$$C<<
E(AE##E(rH   r   c           	     :  K   |dk    rdnd}| d| d}	||||d}
t          j                    }|                                }d}	 	 |                     d	|	|
t          
           d{V  dS # t
          $ r}t          |          }t           |vr |                                |z
  }||k    rt          d|dd| d|           ||dz  }t          	                    d| j
        |||           t          j        t                     d{V  Y d}~nd}~ww xY w)z;Call ``upload_part_finish``, retrying on biz_code 40093001.ru   r   r   r   z/upload_part_finish)rK   rX   rH   r   r   Tr   r   Nz4upload_part_finish persistent retry timed out after z.0fzs (z retries): r   z?[%s] part_finish retryable error, attempt %d, elapsed=%.1fs: %s)r   r   timerm   r   r   r   _BIZ_CODE_PART_RETRYABLEr   r   ro   r   _PART_FINISH_RETRY_INTERVAL)r   rp   rq   rK   rX   rH   r   rP   r   r   r   loopstartr   r   r   elapseds                    r!   r   z'ChunkedUploader._part_finish_with_retry  s      (500{{l77777"$$	
 
 '))			AA''Dt5H (           A A Ac((.0??))++-m++&Q!.PQ Q7>Q QKNQ Q   1(M7GS  
 m$?@@@@@@@@@@@@@@!A	As   $A- -
D7BDDc           	       K   |dk    rdnd}| d| d}d|i}d}t          t          dz             D ]}	 |                     d	||t          
           d{V c S # t          $ rd}	|	}|t          k     rMt
          d|z  z  }
t                              d| j        |dz   |
|	           t          j
        |
           d{V  Y d}	~	d}	~	ww xY wt          dt          dz    d|           )zCall ``complete_upload`` with retry.

        This reuses the ``/files`` endpoint (same as the simple URL-based upload)
        but signals the chunked-completion path by sending only ``upload_id``.
        ru   r   r   r   z/filesrK   Nr   r   r   r   z:[%s] complete_upload attempt %d failed, retry in %.1fs: %szcomplete_upload failed after r   )r   _COMPLETE_UPLOAD_MAX_RETRIESrm   r   r   _COMPLETE_UPLOAD_BASE_DELAYr   r   ro   r   r   r   )r   rp   rq   rK   r   r   r   r   r   r   r   s              r!   r   zChunkedUploader._complete  s      (500{{l*****Y'(,9A=>> 	/ 	/G/!..Dt5H /             	/ 	/ 	/99971<HENN-w{E3  
 "-.........	/ G+a/G G<DG G
 
 	
s   #A
C	%ACC	N)rf   )rg   rh   ri   rj   rk   r   r   r   )rp   r   rq   r   rr   r   rs   r   r   r   r   rS   )rp   r   rq   r   rs   r   r   r   r   r   r   r   r   rJ   )rp   r   rq   r   rr   r   r   r   rK   r   rx   r   ry   rE   rP   rO   rz   r=   r   r   )
rY   r   rU   r   rX   r   r>   r   r   r   )rp   r   rq   r   rK   r   rX   r   rH   r   r   r   rP   rO   r   r   )rp   r   rq   r   rK   r   r   rS   )r+   r,   r-   r.   r   r   r   r}   r   r   r   rC   r"   r!   re   re      s          	         SM SM SM SMr, , , ,H-
 -
 -
 -
^0
 0
 0
 0
d-A -A -A -Af"
 "
 "
 "
 "
 "
r"   re   
size_bytesr   r   c                `    t          |           }dD ]}|dk     r
|dd| c S |dz  }|ddS )z>Return a human-readable file size string (e.g. ``'12.3 MB'``).)BKBMBGBg      @z.1f z TB)rO   )r   sizeunits      r!   r&   r&     sa    D'  &==''''''''r"   rr   r   r   r   c                &   t          | d          5 }|                    |           |                    |          }t          |          |k    r)t	          d|  d| d| dt          |           d	          |cddd           S # 1 swxY w Y   dS )zRead *length* bytes from *file_path* starting at *offset*.

    :raises IOError: If fewer bytes were read than expected (truncated file).
    rbzShort read from z: expected z bytes at offset z, got z (file may be truncated)N)openseekreadr   IOError)rr   r   r   fhrU   s        r!   r   r     s   
 
i		 "
wwvt99L9 L L L L L L(+D		L L L                    s   A(BB
B
r   r   c                n   t          j                    }t          j                    }t          j                    }|t          k    }d}t	          | d          5 }	 |                    d          }|snl|                    |           |                    |           |r-t          |z
  }	|	dk    r|                    |d|	                    |t          |          z  }	 ddd           n# 1 swxY w Y   |                                }
|
|                                |r|                                n|
dS )z0Compute md5, sha1, and md5_10m in a single pass.r   r   Ti   N)r   r   r   )	r   r   r   _MD5_10M_SIZEr   r   updater   r   )rr   r   r   r   r   need_10m
bytes_readr   chunk	remainingfull_md5s              r!   r   r   /  sn   
+--C<>>DkmmG=(HJ	i		 %"
	%GGENNE JJuKK 6)J6	q==NN5)#4555#e**$J
	% 	% % % % % % % % % % % % % % % }}H  *2@7$$&&&	  s   BC**C.1C.r   #List[Callable[[], Awaitable[None]]]rM   r   c                   K   |dk     rd}t          j        |          d	fdt          j        fd| D               d{V  dS )
z=Run a list of thunks with a bounded number in flight at once.r   thunkCallable[[], Awaitable[None]]r   r   c                   K   4 d {V   |              d {V  d d d           d {V  d S # 1 d {V swxY w Y   d S r$   rC   )r   sems    r!   _wrapz$_run_with_concurrency.<locals>._wrapW  s       	 	 	 	 	 	 	 	%''MMMMMMM	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	s   0
::c              3  .   K   | ]} |          V  d S r$   rC   )r~   tr   s     r!   	<genexpr>z(_run_with_concurrency.<locals>.<genexpr>[  s+      335588333333r"   N)r   r   r   r   )r   	Semaphoregather)r   rM   r   r   s     @@r!   r   r   N  s      
 Q

K
(
(C      .3333U333
4444444444r"   )rR   rS   r   rJ   )r   r   r   r   )rr   r   r   r   r   r   r   r   )rr   r   r   r   r   r   )r   r   rM   r   r   r   )2r.   
__future__r   r   r{   r   loggingdataclassesr   r   pathlibr   typingr   r   r	   r
   r   r   !gateway.platforms.qqbot.constantsr   	getLoggerr+   r   r   r   rQ   r   r   r   r   r   r   r   r   r   r   r   r3   r=   rE   rJ   rc   r   rh   re   r&   r   r   r   rC   r"   r!   <module>r      s   B # " " " " "        ( ( ( ( ( ( ( (       A A A A A A A A A A A A A A A A A A A A A A		8	$	$ ! #      ! $     !  
+ + + + +I + + +(P P P P Pi P P PB                        $ $ $ $R YtCH~667H
 H
 H
 H
 H
 H
 H
 H
Z
          >5 5 5 5 5 5r"   