from .protocol import ChunkrClientProtocol
from .config import Configuration, Status, OutputResponse
import asyncio
from datetime import datetime
from pydantic import BaseModel, PrivateAttr
import time
from typing import Optional, Union

class TaskResponse(BaseModel):
    configuration: Configuration
    created_at: datetime
    expires_at: Optional[datetime]
    file_name: Optional[str]
    finished_at: Optional[datetime]
    input_file_url: Optional[str]
    message: str
    output: Optional[OutputResponse]
    page_count: Optional[int]
    pdf_url: Optional[str]
    status: Status
    task_id: str
    task_url: Optional[str]
    _client: Optional[Union[ChunkrClientProtocol]] = PrivateAttr(default=None)

    def with_client(self, client: Union[ChunkrClientProtocol]) -> 'TaskResponse':
        self._client = client
        return self
    
    def _poll_request_sync(self) -> dict:
        """Helper method to make polling request with retry logic (synchronous)"""
        if not self.task_url:
            raise ValueError("Task URL not found in response")

        while True:
            try:
                r = self._client._session.get(self.task_url, headers=self._client._headers())
                r.raise_for_status()
                return r.json()
            except (ConnectionError, TimeoutError) as _:
                print("Connection error while polling the task, retrying...")
                time.sleep(0.5)
            except Exception as e:
                raise

    async def _poll_request_async(self) -> dict:
        """Helper method to make polling request with retry logic (asynchronous)"""
        if not self.task_url:
            raise ValueError("Task URL not found in response")

        while True:
            try:
                r = await self._client._client.get(self.task_url, headers=self._client._headers())
                await r.raise_for_status()
                return await r.json()
            except (ConnectionError, TimeoutError) as _:
                print("Connection error while polling the task, retrying...")
                await asyncio.sleep(0.5)
            except Exception as e:
                raise

    def _check_status(self) -> Optional['TaskResponse']:
        """Helper method to check task status and handle completion/failure"""
        if self.status == "Failed":
            raise ValueError(self.message)
        if self.status not in ("Starting", "Processing"):
            return self
        return None

    def poll(self) -> 'TaskResponse':
        """Poll the task for completion."""
        while True:
            response = self._poll_request_sync()
            self.__dict__.update(response)
            
            if result := self._check_status():
                return result
            
            time.sleep(0.5)

    async def poll_async(self) -> 'TaskResponse':
        """Poll the task for completion asynchronously."""
        while True:
            response = await self._poll_request_async()
            self.__dict__.update(response)
            
            if result := self._check_status():
                return result
            
            await asyncio.sleep(0.5)

    def _get_content(self, content_type: str) -> str:
        """Helper method to get either HTML, Markdown, or raw content."""
        if not self.output:
            return ""
        parts = []
        for c in self.output.chunks:
            for s in c.segments:
                content = getattr(s, content_type)
                if content:
                    parts.append(content)
        return "\n".join(parts)

    def html(self) -> str:
        """Get full HTML for the task"""
        return self._get_content("html")

    def markdown(self) -> str:
        """Get full markdown for the task"""
        return self._get_content("markdown")
        
    def content(self) -> str:
        """Get full text for the task"""
        return self._get_content("content")

class TaskPayload(BaseModel):
    current_configuration: Configuration
    file_name: str
    image_folder_location: str
    input_location: str
    output_location: str
    pdf_location: str
    previous_configuration: Optional[Configuration]
    task_id: str
    user_id: str