# ankur_scraper/output_handler.py

import json
import csv
from pathlib import Path
from datetime import datetime
from ankur_scraper.logging_config import get_logger

info_logger = get_logger("info")
error_logger = get_logger("error")


def save_to_json(data: dict, output_path: str = None) -> str:
    """
    Save scraped data to JSON file.
    
    Args:
        data: Dictionary with 'data' and 'summary' keys
        output_path: Optional custom output path
    
    Returns:
        Path to saved file
    """
    if output_path is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = f"scraped_data_{timestamp}.json"
    
    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
        
        info_logger.info(f"[green]✅ Saved JSON output to:[/] {output_path}")
        return output_path
    except Exception as e:
        error_logger.error(f"[red]Failed to save JSON:[/] {e}")
        return None


def save_to_csv(data: dict, output_path: str = None) -> str:
    """
    Save scraped data to CSV file.
    
    Args:
        data: Dictionary with 'data' and 'summary' keys
        output_path: Optional custom output path
    
    Returns:
        Path to saved file
    """
    if output_path is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = f"scraped_data_{timestamp}.csv"
    
    try:
        with open(output_path, 'w', newline='', encoding='utf-8') as f:
            if not data.get('data'):
                error_logger.warning("No data to save to CSV")
                return None
            
            # Get fieldnames from first item
            fieldnames = ['section', 'source_url', 'extraction_time', 'content']
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            
            for item in data['data']:
                row = {
                    'section': item['metadata']['section'],
                    'source_url': item['metadata']['source_url'],
                    'extraction_time': item['metadata']['extraction_time'],
                    'content': item['content']
                }
                writer.writerow(row)
        
        info_logger.info(f"[green]✅ Saved CSV output to:[/] {output_path}")
        return output_path
    except Exception as e:
        error_logger.error(f"[red]Failed to save CSV:[/] {e}")
        return None


def save_to_txt(data: dict, output_path: str = None) -> str:
    """
    Save scraped data to plain text file.
    
    Args:
        data: Dictionary with 'data' and 'summary' keys
        output_path: Optional custom output path
    
    Returns:
        Path to saved file
    """
    if output_path is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = f"scraped_data_{timestamp}.txt"
    
    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            # Write header
            f.write("=" * 80 + "\n")
            f.write("SCRAPED DATA\n")
            f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write("=" * 80 + "\n\n")
            
            # Write summary
            summary = data.get('summary', {})
            f.write("SUMMARY:\n")
            f.write(f"  Successful pages: {summary.get('successful_pages', 0)}\n")
            f.write(f"  Failed pages: {summary.get('failed_pages', 0)}\n")
            f.write(f"  Total sections: {summary.get('total_sections', 0)}\n")
            f.write(f"  Total pages found: {summary.get('total_pages_found', 0)}\n")
            f.write("\n" + "=" * 80 + "\n\n")
            
            # Write content
            for idx, item in enumerate(data.get('data', []), 1):
                metadata = item['metadata']
                f.write(f"[{idx}] {metadata['section']}\n")
                f.write(f"Source: {metadata['source_url']}\n")
                f.write(f"Time: {metadata['extraction_time']}\n")
                f.write("-" * 80 + "\n")
                f.write(item['content'] + "\n")
                f.write("\n" + "=" * 80 + "\n\n")
        
        info_logger.info(f"[green]✅ Saved TXT output to:[/] {output_path}")
        return output_path
    except Exception as e:
        error_logger.error(f"[red]Failed to save TXT:[/] {e}")
        return None


def save_results(data: dict, output_format: str = "json", output_path: str = None) -> str:
    """
    Save scraped results in the specified format.
    
    Args:
        data: Dictionary with 'data' and 'summary' keys
        output_format: Format to save ('json', 'csv', 'txt', 'all')
        output_path: Optional custom output path (without extension)
    
    Returns:
        Path(s) to saved file(s)
    """
    if not data or not data.get('data'):
        error_logger.warning("No data to save")
        return None
    
    saved_files = []
    
    if output_format == "all":
        # Save in all formats
        for fmt in ['json', 'csv', 'txt']:
            if output_path:
                path = f"{output_path}.{fmt}"
            else:
                path = None
            
            if fmt == 'json':
                result = save_to_json(data, path)
            elif fmt == 'csv':
                result = save_to_csv(data, path)
            elif fmt == 'txt':
                result = save_to_txt(data, path)
            
            if result:
                saved_files.append(result)
    else:
        # Save in specified format
        if output_path and not output_path.endswith(f'.{output_format}'):
            output_path = f"{output_path}.{output_format}"
        
        if output_format == 'json':
            result = save_to_json(data, output_path)
        elif output_format == 'csv':
            result = save_to_csv(data, output_path)
        elif output_format == 'txt':
            result = save_to_txt(data, output_path)
        else:
            error_logger.error(f"Unknown output format: {output_format}")
            return None
        
        if result:
            saved_files.append(result)
    
    return saved_files if len(saved_files) > 1 else (saved_files[0] if saved_files else None)