import csv
import json
import logging
from typing import Any, Dict, List

class ResearchAgent:
    """
    Collects, augments, audits, and outputs lead and signal data for outreach,
    with robust error/memory handling and dynamic variable mapping.
    """
    def __init__(self, db: Any = None, schema: Dict[str, Any] = None):
        self.db = db
        self.schema = schema or {}
        self.memory = []
        self.errors = []
        self.profiles = []
        self.refined = []
        self.mcp_updates = []
        logging.basicConfig(level=logging.INFO)

    def map_raw_data(self, raw_lead: Dict[str, Any], signals: Dict[str, Any], dynamic_vars: Dict[str, Any]) -> Dict[str, Any]:
        mapped = {}
        try:
            for field, rule in (self.schema.get('lead', {})).items():
                mapped[field] = raw_lead.get(field, rule.get('default'))
            mapped.update(signals)
            mapped.update(dynamic_vars)
        except Exception as e:
            self.errors.append(f"Mapping error: {e}")
            logging.error(f"Mapping error: {e}")
        return mapped

    def handle_loop_errors(self, data_list: List[Dict[str, Any]]):
        for idx, entry in enumerate(data_list):
            try:
                assert all(entry.values()), f"Missing value in {entry}"
            except Exception as e:
                self.errors.append(f"Loop error at {idx}: {e}")
                logging.warning(f"Loop error at {idx}: {e}")

    def memory_management(self, current: Dict[str, Any]):
        self.memory.append(current)
        if len(self.memory) > 50:
            dropped = self.memory.pop(0)
            logging.info(f"Dropped old memory: {dropped}")

    def refine_profile(self, profile: Dict[str, Any], messaging_signals: Dict[str, Any]) -> Dict[str, Any]:
        """Output the most outreach-ready version of the profile."""
        refined = {**profile, **messaging_signals}
        self.refined.append(refined)
        return refined

    def handoff(self, profiles: List[Dict[str, Any]], target: str = 'csv', path: str = 'handoff.csv'):
        """Wire/export to csv/mcp/db."""
        try:
            if target == 'csv':
                with open(path, 'w', newline='') as f:
                    writer = csv.DictWriter(f, fieldnames=profiles[0].keys())
                    writer.writeheader()
                    writer.writerows(profiles)
            elif target == 'mcp' and self.db:
                for p in profiles:
                    self.db.update_profile(p)
            elif target == 'db' and self.db:
                self.db.bulk_insert(profiles)
            self.mcp_updates.extend(profiles)
            logging.info(f"Profiles handed off to {target}")
        except Exception as e:
            self.errors.append(f"Handoff error: {e}")
            logging.error(f"Handoff error: {e}")

    def audit(self, profiles: List[Dict[str, Any]]):
        """Audit field completeness, error/correction, and handoff readiness."""
        required_fields = self.schema.get('required', set())
        for idx, profile in enumerate(profiles):
            missing = [f for f in required_fields if not profile.get(f)]
            if missing:
                self.errors.append(f"Profile {idx} missing: {missing}")
                logging.warning(f"Profile {idx} missing: {missing}")

    def export_refined_for_outreach(self, out_path: str = 'refined_for_outreach.csv'):
        try:
            with open(out_path, 'w', newline='') as f:
                writer = csv.DictWriter(f, fieldnames=self.refined[0].keys())
                writer.writeheader()
                writer.writerows(self.refined)
        except Exception as e:
            self.errors.append(f"Export error: {e}")
            logging.error(f"Export error: {e}")

    def get_messaging_recommendations(self, profile: Dict[str, Any]) -> Dict[str, Any]:
        """Stub: Replace with your ML/model recommendation logic."""
        recs = {"recommendation": f"Personalize outreach for {profile.get('name', 'Unknown')}"}
        return recs

    def run_pipeline(self, leads, signals_list, dynamic_vars_list):
        for raw_lead, signals, dynamic_vars in zip(leads, signals_list, dynamic_vars_list):
            mapped = self.map_raw_data(raw_lead, signals, dynamic_vars)
            self.memory_management(mapped)
            self.profiles.append(mapped)
            messaging = self.get_messaging_recommendations(mapped)
            refined = self.refine_profile(mapped, messaging)
        self.audit(self.profiles)
        self.handoff(self.refined)