import csv
import json
from pathlib import Path
from datetime import datetime
from fsm_crawl.blocking.labeler import Labeler
from fsm_crawl.monitoring.callback_pipeline import GenericCallbackPipeline

class RequestResponseLoggingPipeline(GenericCallbackPipeline):
    def __init__(
        self,
        labeler: Labeler,
        output_dir: str = "crawl_logs",
        filename_prefix: str = "requests",
    ):
        self.labeler = labeler
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.filename_prefix = filename_prefix
        self.split_index = 0
        self.request_file = None
        self.response_file = None
        self.request_writer = None
        self.response_writer = None

        self._open_new_files()

    # -------------------------
    # File handling
    # -------------------------
    def _open_new_files(self):
        if self.request_file:
            self.request_file.close()
        if self.response_file:
            self.response_file.close()

        ts = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
        
        # Request file
        request_filename = f"request_{self.filename_prefix}_{self.split_index}_{ts}.csv"
        request_path = self.output_dir / request_filename
        self.request_file = open(request_path, "w", newline="", encoding="utf-8")
        self.request_writer = csv.DictWriter(
            self.request_file,
            fieldnames=[
                "timestamp",
                "request_method",
                "request_resource_type",
                "request_url",
                "request_frame_url",
                "request_headers",
                "label",
            ],
        )
        self.request_writer.writeheader()

        # Response file
        response_filename = f"response_{self.filename_prefix}_{self.split_index}_{ts}.csv"
        response_path = self.output_dir / response_filename
        self.response_file = open(response_path, "w", newline="", encoding="utf-8")
        self.response_writer = csv.DictWriter(
            self.response_file,
            fieldnames=[
                "timestamp",
                "request_method",
                "request_resource_type",
                "request_url",
                "request_frame_url",
                "request_headers",
                "label",
                "response_url",
                "response_status",
                "response_status_text",
                "response_headers",
                "response_from_cache",
                "response_from_service_worker",
            ],
        )
        self.response_writer.writeheader()

    def on_response(self, response):
        try:
            request = response.request

            # Extract request data
            req_data = {
                "timestamp": datetime.utcnow().isoformat(),
                "request_method": request.method,
                "request_resource_type": request.resource_type,
                "request_url": request.url,
                "request_frame_url": request.frame.url if request.frame else None,
                "request_headers": json.dumps(request.headers, ensure_ascii=False),
                "label": self.labeler.label(request.url),
            }

            resp_data = {
                "response_url": response.url,
                "response_status": response.status,
                "response_status_text": response.status_text,
                "response_headers": json.dumps(response.headers, ensure_ascii=False),
                "response_from_cache": response.from_cache,
                "response_from_service_worker": response.from_service_worker,
            }

            # Merge and write to response file
            row = {**req_data, **resp_data}
            self.response_writer.writerow(row)
            self.response_file.flush()

        except Exception as e:
            print(f"[Pipeline] Error logging response: {e}")
            
    async def on_request(self, route, request):
        try:
            # Extract request data
            req_data = {
                "timestamp": datetime.utcnow().isoformat(),
                "request_method": request.method,
                "request_resource_type": request.resource_type,
                "request_url": request.url,
                "request_frame_url": request.frame.url if request.frame else None,
                "request_headers": json.dumps(request.headers, ensure_ascii=False),
                "label": self.labeler.label(request.url),
            }

            # Write to request file
            self.request_writer.writerow(req_data)
            self.request_file.flush()

        except Exception as e:
            print(f"[Pipeline] Error logging request: {e}")
        
        await route.continue_()

    async def on_iteration_split(self):
        self.split_index += 1
        self._open_new_files()

    async def on_crawl_end(self):
        if self.request_file:
            self.request_file.close()
        if self.response_file:
            self.response_file.close()
