# -*- coding: utf-8 -*- 
import pytest
from iterable.datatypes import WARCIterable
from iterable.codecs import GZIPCodec


class TestWARC:
    def test_id(self):
        datatype_id = WARCIterable.id()
        assert datatype_id == 'warc'

    def test_flatonly(self):
        flag = WARCIterable.is_flatonly()
        assert flag == False

    def test_has_totals(self):
        flag = WARCIterable.has_totals()
        assert flag == False

    def test_openclose(self):
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj)
        iterable.close()

    def test_read_one(self):
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj)
        record = iterable.read()
        # Check that record has expected structure
        assert isinstance(record, dict)
        assert 'rec_type' in record
        assert 'rec_headers' in record
        assert 'content_type' in record
        assert 'length' in record
        iterable.close()

    def test_reset(self):
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj)
        record1 = iterable.read()
        iterable.reset()
        record2 = iterable.read()
        # After reset, should read the first record again
        assert record1['rec_type'] == record2['rec_type']
        assert record1.get('target_uri') == record2.get('target_uri')
        iterable.close()

    def test_next(self):
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj)
        record1 = next(iterable)
        assert isinstance(record1, dict)
        assert 'rec_type' in record1
        iterable.reset()
        record2 = next(iterable)
        assert record1['rec_type'] == record2['rec_type']
        iterable.close()

    def test_count(self):
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj)
        n = 0
        for record in iterable:
            assert isinstance(record, dict)
            assert 'rec_type' in record
            n += 1
        # Should have at least one record
        assert n > 0
        iterable.close()

    def test_iterate_all(self):
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj)
        records = []
        for record in iterable:
            assert isinstance(record, dict)
            assert 'rec_type' in record
            assert 'rec_headers' in record
            records.append(record)
        # Should have at least one record
        assert len(records) > 0
        iterable.close()

    def test_read_bulk(self):
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj)
        chunk = iterable.read_bulk(3)
        assert isinstance(chunk, list)
        assert len(chunk) > 0
        for record in chunk:
            assert isinstance(record, dict)
            assert 'rec_type' in record
        iterable.close()

    def test_stop_iteration_after_exhaustion(self):
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj)
        # Read all records
        try:
            while True:
                _ = iterable.read()
        except StopIteration:
            pass
        # Should raise StopIteration on next read
        with pytest.raises(StopIteration):
            _ = iterable.read()
        iterable.close()

    def test_record_structure(self):
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj)
        record = iterable.read()
        # Check required fields
        assert 'rec_type' in record
        assert 'rec_headers' in record
        assert isinstance(record['rec_headers'], dict)
        assert 'content_type' in record
        assert 'length' in record
        assert isinstance(record['length'], int)
        iterable.close()

    def test_totals_not_supported(self):
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj)
        with pytest.raises(NotImplementedError):
            iterable.totals()
        iterable.close()

    def test_write_single_record(self):
        """Test writing a single WARC record"""
        import os
        import tempfile
        
        test_file = 'testdata/test_warc_write.warc'
        os.makedirs('testdata', exist_ok=True)
        
        # Write a record
        iterable = WARCIterable(test_file, mode='w')
        record = {
            'rec_type': 'response',
            'target_uri': 'http://example.com/test',
            'rec_headers': {
                'WARC-Date': '2024-01-01T00:00:00Z',
                'WARC-Record-ID': '<urn:uuid:test-123>',
            },
            'content': b'Test content'
        }
        iterable.write(record)
        iterable.close()
        
        # Verify file was created
        assert os.path.exists(test_file)
        assert os.path.getsize(test_file) > 0
        
        # Cleanup
        if os.path.exists(test_file):
            os.unlink(test_file)

    def test_write_bulk_records(self):
        """Test writing multiple WARC records"""
        import os
        import tempfile
        
        test_file = 'testdata/test_warc_write_bulk.warc'
        os.makedirs('testdata', exist_ok=True)
        
        # Write multiple records
        iterable = WARCIterable(test_file, mode='w')
        records = [
            {
                'rec_type': 'response',
                'target_uri': 'http://example.com/test1',
                'content': b'Content 1'
            },
            {
                'rec_type': 'response',
                'target_uri': 'http://example.com/test2',
                'content': b'Content 2'
            }
        ]
        iterable.write_bulk(records)
        iterable.close()
        
        # Verify file was created
        assert os.path.exists(test_file)
        assert os.path.getsize(test_file) > 0
        
        # Cleanup
        if os.path.exists(test_file):
            os.unlink(test_file)

    def test_write_read_roundtrip(self):
        """Test writing and then reading back WARC records"""
        import os
        
        test_file = 'testdata/test_warc_roundtrip.warc'
        os.makedirs('testdata', exist_ok=True)
        
        # Write records
        writer = WARCIterable(test_file, mode='w')
        test_records = [
            {
                'rec_type': 'response',
                'target_uri': 'http://example.com/test1',
                'rec_headers': {
                    'WARC-Date': '2024-01-01T00:00:00Z',
                },
                'content': b'Test content 1'
            },
            {
                'rec_type': 'response',
                'target_uri': 'http://example.com/test2',
                'content': b'Test content 2'
            }
        ]
        writer.write_bulk(test_records)
        writer.close()
        
        # Read back
        reader = WARCIterable(test_file, mode='r')
        read_records = []
        for record in reader:
            read_records.append(record)
        reader.close()
        
        # Verify we can read what we wrote
        assert len(read_records) >= len(test_records)
        # Check that target URIs match
        assert any(r.get('target_uri') == 'http://example.com/test1' for r in read_records)
        assert any(r.get('target_uri') == 'http://example.com/test2' for r in read_records)
        
        # Cleanup
        if os.path.exists(test_file):
            os.unlink(test_file)

    def test_write_with_codec(self):
        """Test writing WARC records with compression codec"""
        import os
        
        test_file = 'testdata/test_warc_write_codec.warc.gz'
        os.makedirs('testdata', exist_ok=True)
        
        # Write with GZIP codec
        codec = GZIPCodec(test_file, mode='w', open_it=True)
        iterable = WARCIterable(codec=codec, mode='w')
        record = {
            'rec_type': 'response',
            'target_uri': 'http://example.com/test',
            'content': b'Compressed content'
        }
        iterable.write(record)
        iterable.close()
        
        # Verify file was created
        assert os.path.exists(test_file)
        assert os.path.getsize(test_file) > 0
        
        # Cleanup
        if os.path.exists(test_file):
            os.unlink(test_file)

    def test_write_error_not_opened_for_writing(self):
        """Test that writing to a read-only iterable raises error"""
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj, mode='r')
        record = {'rec_type': 'response', 'target_uri': 'http://example.com'}
        with pytest.raises(RuntimeError, match="WARC file not opened for writing"):
            iterable.write(record)
        iterable.close()

    def test_write_with_string_content(self):
        """Test writing WARC record with string content (should be encoded)"""
        import os
        
        test_file = 'testdata/test_warc_write_string.warc'
        os.makedirs('testdata', exist_ok=True)
        
        iterable = WARCIterable(test_file, mode='w')
        record = {
            'rec_type': 'response',
            'target_uri': 'http://example.com/test',
            'content': 'String content that should be encoded'
        }
        iterable.write(record)
        iterable.close()
        
        # Verify file was created
        assert os.path.exists(test_file)
        
        # Cleanup
        if os.path.exists(test_file):
            os.unlink(test_file)

    def test_write_with_url_field(self):
        """Test writing WARC record using 'url' field instead of 'target_uri'"""
        import os
        
        test_file = 'testdata/test_warc_write_url.warc'
        os.makedirs('testdata', exist_ok=True)
        
        iterable = WARCIterable(test_file, mode='w')
        record = {
            'rec_type': 'response',
            'url': 'http://example.com/from-url-field',
            'content': b'Content'
        }
        iterable.write(record)
        iterable.close()
        
        # Verify file was created
        assert os.path.exists(test_file)
        
        # Cleanup
        if os.path.exists(test_file):
            os.unlink(test_file)

    def test_import_error_when_warcio_missing(self):
        """Test that ImportError is raised when warcio is not available"""
        # This test would require mocking the import, which is complex
        # The actual ImportError is raised in __init__ if HAS_WARCIO is False
        # Since warcio is likely installed in test environment, we verify the check exists
        try:
            from iterable.datatypes.warc import HAS_WARCIO
            # If we can import, verify the check is in place
            assert isinstance(HAS_WARCIO, bool)
        except ImportError:
            # If warcio module itself can't be imported, that's expected
            pass

    def test_http_headers_extraction(self):
        """Test that HTTP headers and metadata are extracted from WARC records"""
        codecobj = GZIPCodec('fixtures/sample.warc.gz', mode='r', open_it=True)
        iterable = WARCIterable(codec=codecobj)
        
        # Read all records and check for HTTP headers in response/request records
        found_http_headers = False
        found_response_metadata = False
        
        for record in iterable:
            # Check if record has HTTP headers (for response/request records)
            if record.get('rec_type') in ['response', 'request']:
                # HTTP headers should be extracted if available
                if 'http_headers' in record:
                    found_http_headers = True
                    # http_headers should be a dict (or None)
                    assert record['http_headers'] is None or isinstance(record['http_headers'], dict)
                
                # For response records, check for status code and status line
                if record.get('rec_type') == 'response':
                    if 'http_status_code' in record:
                        found_response_metadata = True
                        assert isinstance(record['http_status_code'], (int, str))
                    if 'http_status_line' in record:
                        found_response_metadata = True
                        assert isinstance(record['http_status_line'], str)
                
                # For request records, check for request line and parsed components
                elif record.get('rec_type') == 'request':
                    if 'http_request_line' in record:
                        found_response_metadata = True
                        assert isinstance(record['http_request_line'], str)
                    if 'http_method' in record:
                        assert isinstance(record['http_method'], str)
                    if 'http_path' in record:
                        assert isinstance(record['http_path'], str)
                    if 'http_protocol' in record:
                        assert isinstance(record['http_protocol'], str)
            
            # Limit iteration to avoid processing entire large file
            if found_http_headers and found_response_metadata:
                break
        
        iterable.close()
        
        # Note: We don't assert that http_headers must be found, as the test fixture
        # may not contain HTTP response/request records, but if they exist, they should be extracted

