import time
from pathlib import Path
from typing import Optional, Callable, Dict, Any, List
import httpx

from agent.config import AgentConfig, load_config
from agent.chunker import FileChunker
from agent.state_db import state_db

class ChunkUploader:
    """HTTP/HTTPS Chunk transfer client with automatic resumption and retry backoff."""

    def __init__(self, config: Optional[AgentConfig] = None):
        self.config = config or load_config()

    def _get_headers(self) -> Dict[str, str]:
        if not self.config.device_id or not self.config.device_token:
            raise ValueError("Agent is not registered. Run 'agent_cli.py register' first.")
        return {
            "X-Device-Id": self.config.device_id,
            "X-Device-Token": self.config.device_token
        }

    def upload_file(
        self,
        filepath: Path,
        job_id: Optional[int] = None,
        progress_callback: Optional[Callable[[int, int, float], None]] = None,
        max_retries: int = 5
    ) -> Dict[str, Any]:
        """
        Uploads a file chunk by chunk. If interrupted, querying the server will return
        already received chunks, allowing immediate resumption without re-uploading completed parts.
        """
        filepath = Path(filepath).resolve()
        chunker = FileChunker(filepath=filepath, chunk_size=self.config.chunk_size)
        full_sha256 = chunker.full_sha256
        filename = filepath.name

        headers = self._get_headers()
        base_url = self.config.server_url.rstrip("/")

        with httpx.Client(base_url=base_url, headers=headers, timeout=60.0) as client:
            # 1. Initialize or resume upload session
            init_payload = {
                "filename": filename,
                "file_size": chunker.file_size,
                "sha256": full_sha256,
                "chunk_size": chunker.chunk_size,
                "job_id": job_id
            }

            resp = client.post("/api/upload/session", json=init_payload)
            resp.raise_for_status()
            session_data = resp.json()

            session_code = session_data["session_code"]
            total_chunks = session_data["total_chunks"]
            received_chunks = set(session_data.get("received_chunks", []))

            state_db.save_session(
                filepath=str(filepath),
                sha256=full_sha256,
                session_code=session_code,
                total_chunks=total_chunks,
                chunk_size=chunker.chunk_size,
                status="UPLOADING"
            )

            # 2. Upload missing chunks
            missing_chunks = [i for i in range(total_chunks) if i not in received_chunks]

            for idx in missing_chunks:
                chunk_bytes, chunk_hash = chunker.get_chunk(idx)
                
                # Retry loop with exponential backoff for network resilience
                success = False
                attempt = 0
                while not success and attempt < max_retries:
                    try:
                        chunk_headers = {
                            "Content-Type": "application/octet-stream",
                            "X-Chunk-Index": str(idx),
                            "X-Chunk-SHA256": chunk_hash
                        }
                        c_resp = client.post(
                            f"/api/upload/{session_code}/chunk?chunk_index={idx}&chunk_sha256={chunk_hash}",
                            content=chunk_bytes,
                            headers=chunk_headers
                        )
                        c_resp.raise_for_status()
                        success = True
                    except Exception as ex:
                        attempt += 1
                        if attempt >= max_retries:
                            raise ConnectionError(f"Failed to upload chunk {idx} after {max_retries} attempts: {str(ex)}")
                        time.sleep(2 ** attempt)

                received_chunks.add(idx)
                if progress_callback:
                    pct = round((len(received_chunks) / total_chunks) * 100, 2)
                    progress_callback(len(received_chunks), total_chunks, pct)

            # 3. Complete and verify full file assembly on server
            complete_resp = client.post(f"/api/upload/{session_code}/complete")
            complete_resp.raise_for_status()
            result = complete_resp.json()

            # 4. Update local state database
            state_db.mark_session_completed(
                filepath=str(filepath),
                sha256=full_sha256,
                file_size=chunker.file_size,
                job_id=job_id
            )

            return result