from __future__ import annotations

import time
import os
import csv
import io
import requests
from typing import Dict, Any, Union, Optional
from pathlib import Path

from ..types import (
    CreateImportParameters,
    CreateImportResponse,
    Import,
    ListImportsResponse,
    UpdateImport,
)
from ..core.base import WebsetsBaseClient

class ImportsClient(WebsetsBaseClient):
    """Client for managing Imports."""
    
    def __init__(self, client):
        super().__init__(client)

    def create(
        self, 
        params: Union[Dict[str, Any], CreateImportParameters], 
        *,
        csv_data: Optional[Union[str, Path]] = None
    ) -> Union[CreateImportResponse, Import]:
        """Create a new import to upload your data into Websets.
        
        Imports can be used to:
        - Enrich: Enhance your data with additional information using our AI-powered enrichment engine
        - Search: Query your data using Websets' agentic search with natural language filters
        - Exclude: Prevent duplicate or already known results from appearing in your searches

        Args:
            params (CreateImportParameters): The parameters for creating an import.
            csv_data (Union[str, Path], optional): CSV data to upload. Can be raw CSV string or file path.
                                                   When provided, size and count will be automatically calculated 
                                                   if not specified in params.
        
        Returns:
            CreateImportResponse: If csv_data is None (traditional usage with upload URL).
            Import: If csv_data is provided (uploaded and processing).
        """
        if csv_data is not None:
            if isinstance(csv_data, (str, Path)) and os.path.isfile(csv_data):
                with open(csv_data, 'r', encoding='utf-8') as f:
                    csv_content = f.read()
            else:
                csv_content = str(csv_data)
            
            if isinstance(params, dict):
                current_size = params.get('size')
                current_count = params.get('count')
            else:
                current_size = getattr(params, 'size', None)
                current_count = getattr(params, 'count', None)
            
            if current_size is None or current_count is None:
                calculated_size = len(csv_content.encode('utf-8'))
                csv_reader = csv.reader(io.StringIO(csv_content))
                rows = list(csv_reader)
                calculated_count = max(0, len(rows) - 1)
                
                if isinstance(params, dict):
                    params = params.copy()
                    if current_size is None:
                        params['size'] = calculated_size
                    if current_count is None:
                        params['count'] = calculated_count
                else:
                    params_dict = params.model_dump()
                    if current_size is None:
                        params_dict['size'] = calculated_size
                    if current_count is None:
                        params_dict['count'] = calculated_count
                    params = CreateImportParameters.model_validate(params_dict)
        
        response = self.request("/v0/imports", data=params)
        import_response = CreateImportResponse.model_validate(response)
        
        if csv_data is None:
            return import_response
        
        upload_response = requests.put(
            import_response.upload_url,
            data=csv_content,
            headers={'Content-Type': 'text/csv'}
        )
        upload_response.raise_for_status()
        
        return self.get(import_response.id)

    def get(self, import_id: str) -> Import:
        """Get a specific import.
        
        Args:
            import_id (str): The id of the Import.
        
        Returns:
            Import: The retrieved import.
        """
        response = self.request(f"/v0/imports/{import_id}", method="GET")
        return Import.model_validate(response)

    def list(self, *, cursor: Optional[str] = None, limit: Optional[int] = None) -> ListImportsResponse:
        """List all imports.
        
        Args:
            cursor (str, optional): The cursor to paginate through the results.
            limit (int, optional): The number of results to return (1-200, default 25).
        
        Returns:
            ListImportsResponse: List of imports with pagination info.
        """
        params = {
            k: v 
            for k, v in {
                "cursor": cursor, 
                "limit": limit
            }.items() 
            if v is not None
        }
        response = self.request("/v0/imports", params=params, method="GET")
        return ListImportsResponse.model_validate(response)

    def update(self, import_id: str, params: Union[Dict[str, Any], UpdateImport]) -> Import:
        """Update an import configuration.
        
        Args:
            import_id (str): The id of the Import.
            params (UpdateImport): The parameters for updating an import.
        
        Returns:
            Import: The updated import.
        """
        response = self.request(f"/v0/imports/{import_id}", data=params, method="PATCH")
        return Import.model_validate(response)

    def delete(self, import_id: str) -> Import:
        """Delete an import.
        
        Args:
            import_id (str): The id of the Import.
        
        Returns:
            Import: The deleted import.
        """
        response = self.request(f"/v0/imports/{import_id}", method="DELETE")
        return Import.model_validate(response)

    def wait_until_completed(self, import_id: str, *, timeout: int = 1800, poll_interval: int = 5) -> Import:
        """Wait until an import is completed or failed.
        
        Args:
            import_id (str): The id of the Import.
            timeout (int, optional): Maximum time to wait in seconds. Defaults to 1800 (30 minutes).
            poll_interval (int, optional): Time to wait between polls in seconds. Defaults to 5.
            
        Returns:
            Import: The import once it's completed or failed.
            
        Raises:
            TimeoutError: If the import does not complete within the timeout period.
        """
        start_time = time.time()
        while True:
            import_obj = self.get(import_id)
            if import_obj.status in ['completed', 'failed']:
                return import_obj
                
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Import {import_id} did not complete within {timeout} seconds")
                
            time.sleep(poll_interval) 