import os
import json
import shutil
import requests
from typing import List, Dict
import argparse

class DataCollector:
    def __init__(self, output_base_dir: str):
        self.output_base_dir = output_base_dir
        os.makedirs(self.output_base_dir, exist_ok=True)

    def collect_from_jsonl(self, file_path: str, fields: List[str]):
        """
        Extracts specific fields from a JSONL file and categorizes them.
        """
        print(f"Collecting from {file_path}...")
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                try:
                    data = json.loads(line)
                    # Create a unified 'text' field for the trainer
                    text_parts = []
                    for field in fields:
                        if field in data:
                            val = data[field]
                            if isinstance(val, list):
                                val = " ".join(val)
                            text_parts.append(val)
                    
                    if text_parts:
                        data["text"] = ". ".join(text_parts)
                    
                    # Use the first available field as category
                    category = fields[0] if fields else "general"
                    self._store_data(data, category)
                except Exception as e:
                    print(f"Error parsing line: {e}")

    def collect_from_url(self, url: str, fields: List[str], file_type: str = 'jsonl'):
        """
        Downloads a file from a URL and processes it.
        """
        print(f"Downloading from {url}...")
        try:
            response = requests.get(url, stream=True)
            response.raise_for_status()
            
            temp_file = os.path.join(self.output_base_dir, "temp_downloaded_data")
            with open(temp_file, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            
            if file_type == 'jsonl':
                self.collect_from_jsonl(temp_file, fields)
            else:
                # Assume raw text or other handling
                self.collect_from_directory(os.path.dirname(temp_file), [os.path.basename(temp_file)], fields[0] if fields else "web")
            
            os.remove(temp_file)
            print("Download and ingestion complete.")
        except Exception as e:
            print(f"Error downloading from URL: {e}")

    def collect_from_directory(self, dir_path: str, extensions: List[str], field_label: str):
        """
        Collects files from a directory and labels them with a field.
        """
        print(f"Scanning directory {dir_path} for {extensions}...")
        for root, _, files in os.walk(dir_path):
            for file in files:
                if any(file.endswith(ext) for ext in extensions):
                    file_path = os.path.join(root, file)
                    # For non-text data, we might just store references or metadata
                    # For now, let's assume we're creating a JSON record for the trainer
                    data = {
                        "file_path": file_path,
                        "field": field_label,
                        "base_name": file
                    }
                    self._store_data(data, field_label)

    def _store_data(self, data: Dict, category: str):
        category_dir = os.path.join(self.output_base_dir, category)
        os.makedirs(category_dir, exist_ok=True)
        
        # Save as individual json for granular control or append to a master for that category
        master_file = os.path.join(category_dir, "collected_data.jsonl")
        with open(master_file, 'a', encoding='utf-8') as f:
            f.write(json.dumps(data) + "\n")

def main():
    parser = argparse.ArgumentParser(description="Data Collector for Fine-tuner")
    parser.add_argument("--source", type=str, required=True, help="Source file, directory, or URL")
    parser.add_argument("--type", choices=['jsonl', 'dir', 'url'], required=True, help="Type of source")
    parser.add_argument("--fields", nargs="+", help="Fields to extract or label")
    parser.add_argument("--output", type=str, default="data/processed", help="Output directory")
    
    args = parser.parse_args()
    
    collector = DataCollector(args.output)
    
    if args.type == 'jsonl':
        collector.collect_from_jsonl(args.source, args.fields or ["text"])
    elif args.type == 'dir':
        collector.collect_from_directory(args.source, [".txt", ".py", ".md"], args.fields[0] if args.fields else "code")
    elif args.type == 'url':
        collector.collect_from_url(args.source, args.fields or ["text"], 'jsonl')

if __name__ == "__main__":
    main()