#!/usr/bin/env python3
"""Prospector: Find leads matching your ICP via Exa + Apollo."""

import csv
import json
import logging
import os
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Optional

import httpx

# ============================================================================
# Configuration
# ============================================================================

CONFIG_PATH = Path.home() / ".config" / "prospector" / "config.json"
DESKTOP_PATH = Path.home() / "Desktop"

# Environment variable overrides (preferred for avoiding plaintext config)
ENV_EXA_KEY = "PROSPECTOR_EXA_API_KEY"
ENV_APOLLO_KEY = "PROSPECTOR_APOLLO_API_KEY"
ENV_ATTIO_KEY = "PROSPECTOR_ATTIO_API_KEY"

logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger(__name__)


@dataclass
class Config:
    """Prospector configuration with API keys."""

    exa_api_key: str
    apollo_api_key: str
    attio_api_key: Optional[str] = None

    @classmethod
    def load(cls) -> "Config":
        """Load config from env or disk. Raises if not found or invalid."""
        env_exa = os.getenv(ENV_EXA_KEY)
        env_apollo = os.getenv(ENV_APOLLO_KEY)
        env_attio = os.getenv(ENV_ATTIO_KEY)

        data: Optional[dict] = None
        if CONFIG_PATH.exists():
            # Verify secure permissions (owner-only)
            mode = CONFIG_PATH.stat().st_mode
            if mode & 0o077:
                raise PermissionError(
                    f"Config has insecure permissions. Run: chmod 600 {CONFIG_PATH}"
                )

            with open(CONFIG_PATH) as f:
                data = json.load(f)

        exa_api_key = env_exa or (data.get("exa_api_key") if data else None)
        apollo_api_key = env_apollo or (data.get("apollo_api_key") if data else None)
        attio_api_key = env_attio or (data.get("attio_api_key") if data else None)

        if not exa_api_key or not apollo_api_key:
            raise FileNotFoundError(
                "Run /prospector:setup first or set "
                f"{ENV_EXA_KEY}/{ENV_APOLLO_KEY} environment variables"
            )

        return cls(
            exa_api_key=exa_api_key,
            apollo_api_key=apollo_api_key,
            attio_api_key=attio_api_key,
        )

    def save(self) -> None:
        """Save config to disk with secure permissions."""
        CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
        with open(CONFIG_PATH, "w") as f:
            json.dump(
                {
                    "exa_api_key": self.exa_api_key,
                    "apollo_api_key": self.apollo_api_key,
                    "attio_api_key": self.attio_api_key,
                },
                f,
                indent=2,
            )
        os.chmod(CONFIG_PATH, 0o600)


@dataclass
class Lead:
    """A single lead with company and contact info."""

    company_name: str
    company_domain: str
    company_size: str
    industry: str
    location: str
    contact_name: str
    contact_title: str
    contact_email: str
    contact_linkedin: str
    source: str = "exa+apollo"


# ============================================================================
# API Key Validation
# ============================================================================


def validate_exa_key(api_key: str) -> bool:
    """Validate Exa API key with a minimal search."""
    try:
        with httpx.Client(timeout=30) as client:
            resp = client.post(
                "https://api.exa.ai/search",
                headers={"x-api-key": api_key, "Content-Type": "application/json"},
                json={"query": "test", "numResults": 1},
            )
            return resp.status_code == 200
    except Exception:
        return False


def validate_apollo_key(api_key: str) -> bool:
    """Validate Apollo API key with a minimal search."""
    try:
        with httpx.Client(timeout=30) as client:
            resp = client.post(
                "https://api.apollo.io/api/v1/mixed_people/search",
                headers={"x-api-key": api_key, "Content-Type": "application/json"},
                json={"per_page": 1},
            )
            return resp.status_code == 200
    except Exception:
        return False


def validate_attio_key(api_key: str) -> bool:
    """Validate Attio API key by fetching workspace info."""
    try:
        with httpx.Client(timeout=30) as client:
            resp = client.get(
                "https://api.attio.com/v2/self",
                headers={"Authorization": f"Bearer {api_key}"},
            )
            return resp.status_code == 200
    except Exception:
        return False


# ============================================================================
# API Clients
# ============================================================================


def search_exa(
    query: str,
    num_results: int,
    api_key: str,
    client: httpx.Client,
) -> list[dict]:
    """Search Exa for companies matching query."""
    logger.info(f"Searching Exa for: {query[:60]}...")

    resp = client.post(
        "https://api.exa.ai/search",
        headers={"x-api-key": api_key, "Content-Type": "application/json"},
        json={
            "query": query,
            "category": "company",
            "numResults": num_results,
            "type": "auto",
        },
    )
    resp.raise_for_status()
    results = resp.json().get("results", [])
    logger.info(f"Found {len(results)} companies")
    return results


def enrich_apollo(
    domain: str,
    api_key: str,
    client: httpx.Client,
    per_company: int = 3,
) -> list[dict]:
    """Find contacts at a company via Apollo."""
    resp = client.post(
        "https://api.apollo.io/api/v1/mixed_people/search",
        headers={"x-api-key": api_key, "Content-Type": "application/json"},
        json={
            "q_organization_domains_list[]": [domain],
            "person_seniorities[]": ["c_suite", "vp", "director", "manager"],
            "per_page": per_company,
        },
    )
    if resp.status_code != 200:
        logger.warning(f"Apollo enrichment failed for {domain}: {resp.status_code}")
        return []
    return resp.json().get("people", [])


def sync_attio(
    leads: list[Lead],
    api_key: str,
    client: httpx.Client,
) -> tuple[int, int]:
    """Sync leads to Attio. Returns (companies_synced, people_synced)."""
    companies_synced = 0
    people_synced = 0
    company_ids: dict[str, str] = {}  # domain -> record_id

    # Upsert companies first
    unique_domains = {lead.company_domain for lead in leads}
    logger.info(f"Syncing {len(unique_domains)} companies to Attio...")

    for domain in unique_domains:
        lead = next(l for l in leads if l.company_domain == domain)
        try:
            resp = client.put(
                "https://api.attio.com/v2/objects/companies/records",
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                params={"matching_attribute": "domains"},
                json={
                    "data": {
                        "values": {
                            "name": [{"value": lead.company_name}],
                            "domains": [{"domain": domain}],
                        }
                    }
                },
            )
            if resp.status_code == 200:
                record_id = resp.json().get("data", {}).get("id", {}).get("record_id")
                if record_id:
                    company_ids[domain] = record_id
                    companies_synced += 1
        except Exception as e:
            logger.warning(f"Failed to sync company {domain}: {e}")

    # Create people linked to companies
    logger.info(f"Syncing {len(leads)} contacts to Attio...")

    for lead in leads:
        if not lead.contact_email:
            continue
        company_record_id = company_ids.get(lead.company_domain)
        if not company_record_id:
            continue

        names = lead.contact_name.split(" ", 1)
        first_name = names[0]
        last_name = names[1] if len(names) > 1 else ""

        try:
            resp = client.post(
                "https://api.attio.com/v2/objects/people/records",
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                json={
                    "data": {
                        "values": {
                            "name": [
                                {
                                    "full_name": lead.contact_name,
                                    "first_name": first_name,
                                    "last_name": last_name,
                                }
                            ],
                            "email_addresses": [{"email_address": lead.contact_email}],
                            "job_title": [{"value": lead.contact_title}],
                            "company": [
                                {
                                    "target_object": "companies",
                                    "target_record_id": company_record_id,
                                }
                            ],
                        }
                    }
                },
            )
            if resp.status_code in (200, 201):
                people_synced += 1
        except Exception as e:
            logger.warning(f"Failed to sync person {lead.contact_email}: {e}")

    return companies_synced, people_synced


# ============================================================================
# Query Building
# ============================================================================


def build_exa_query(icp: dict) -> str:
    """Convert ICP answers to Exa search query."""
    parts = []

    # Industry
    industry = icp.get("industry")
    if industry and industry != "Any":
        parts.append(f"{industry} companies")
    else:
        parts.append("companies")

    # Funding stage
    funding = icp.get("funding_stage")
    if funding and funding != "Any":
        parts.append(f"that have raised {funding} funding")

    # Company size
    size = icp.get("company_size")
    if size and size != "Any":
        size_map = {
            "1-10": "with 1-10 employees",
            "11-50": "with 11-50 employees",
            "51-200": "with 51-200 employees",
            "201-500": "with 201-500 employees",
            "500+": "with over 500 employees",
        }
        size_text = size_map.get(size)
        if size_text:
            parts.append(size_text)

    # Geography
    geo = icp.get("geography")
    if geo and geo != "Any":
        parts.append(f"based in {geo}")

    # Keywords
    keywords = icp.get("keywords")
    if keywords:
        parts.append(f"focused on {keywords}")

    return " ".join(parts)


# ============================================================================
# CSV Export
# ============================================================================


def get_output_path() -> Path:
    """Get CSV output path, handling filename collisions."""
    base = DESKTOP_PATH / f"prospector-leads-{date.today()}.csv"

    if not base.exists():
        return base

    # Handle collisions: -2, -3, etc.
    counter = 2
    while True:
        path = DESKTOP_PATH / f"prospector-leads-{date.today()}-{counter}.csv"
        if not path.exists():
            return path
        counter += 1


def write_csv(leads: list[Lead], path: Path) -> None:
    """Write leads to CSV with UTF-8 BOM for Excel compatibility."""
    with open(path, "w", newline="", encoding="utf-8-sig") as f:
        writer = csv.writer(f)
        writer.writerow(
            [
                "company_name",
                "company_domain",
                "company_size",
                "industry",
                "location",
                "contact_name",
                "contact_title",
                "contact_email",
                "contact_linkedin",
                "source",
            ]
        )
        for lead in leads:
            writer.writerow(
                [
                    lead.company_name,
                    lead.company_domain,
                    lead.company_size,
                    lead.industry,
                    lead.location,
                    lead.contact_name,
                    lead.contact_title,
                    lead.contact_email,
                    lead.contact_linkedin,
                    lead.source,
                ]
            )


# ============================================================================
# Main Entry Points
# ============================================================================


def run_search(icp: dict, num_contacts: int = 50) -> list[Lead]:
    """
    Main search workflow.

    Args:
        icp: Dict with keys: industry, company_size, funding_stage, geography, keywords
        num_contacts: Target number of contacts to find

    Returns:
        List of Lead objects
    """
    config = Config.load()
    leads: list[Lead] = []

    # Build query and search Exa
    query = build_exa_query(icp)

    # Estimate companies needed (3 contacts per company)
    num_companies = max(10, num_contacts // 3)

    with httpx.Client(timeout=60) as client:
        # Search Exa
        companies = search_exa(query, num_companies, config.exa_api_key, client)

        if not companies:
            logger.warning("No companies found matching criteria")
            return []

        # Enrich via Apollo
        logger.info(f"Enriching contacts for {len(companies)} companies...")
        for i, company in enumerate(companies):
            # Extract domain from URL
            url = company.get("url", "")
            domain = url.replace("https://", "").replace("http://", "").split("/")[0]
            if not domain or domain.startswith("linkedin"):
                continue

            contacts = enrich_apollo(domain, config.apollo_api_key, client)

            for contact in contacts:
                leads.append(
                    Lead(
                        company_name=company.get("title", "").split(" - ")[0],
                        company_domain=domain,
                        company_size=icp.get("company_size", ""),
                        industry=icp.get("industry", ""),
                        location=contact.get("city", "") or icp.get("geography", ""),
                        contact_name=f"{contact.get('first_name', '')} {contact.get('last_name', '')}".strip(),
                        contact_title=contact.get("title", ""),
                        contact_email=contact.get("email", ""),
                        contact_linkedin=contact.get("linkedin_url", ""),
                    )
                )

            # Progress
            if (i + 1) % 5 == 0:
                logger.info(
                    f"Processed {i + 1}/{len(companies)} companies, {len(leads)} contacts found"
                )

    logger.info(f"Total: {len(leads)} contacts from {len(companies)} companies")
    return leads


def export_csv(leads: list[Lead]) -> Path:
    """Export leads to CSV on Desktop."""
    path = get_output_path()
    write_csv(leads, path)
    logger.info(f"Saved {len(leads)} leads to {path}")
    return path


def sync_to_attio(leads: list[Lead]) -> tuple[int, int]:
    """Sync leads to Attio CRM."""
    config = Config.load()
    if not config.attio_api_key:
        raise ValueError("Attio not configured. Run /prospector:setup to add Attio key.")

    with httpx.Client(timeout=60) as client:
        return sync_attio(leads, config.attio_api_key, client)


# ============================================================================
# CLI Interface (for direct script usage)
# ============================================================================

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Prospector: Find leads matching your ICP")
    parser.add_argument("--industry", default="Any", help="Target industry")
    parser.add_argument("--size", default="Any", help="Company size")
    parser.add_argument("--funding", default="Any", help="Funding stage")
    parser.add_argument("--geography", default="Any", help="Geography")
    parser.add_argument("--keywords", default="", help="Keywords")
    parser.add_argument("--count", type=int, default=50, help="Number of contacts")
    parser.add_argument("--sync-attio", action="store_true", help="Sync to Attio after search")

    args = parser.parse_args()

    icp = {
        "industry": args.industry,
        "company_size": args.size,
        "funding_stage": args.funding,
        "geography": args.geography,
        "keywords": args.keywords,
    }

    leads = run_search(icp, num_contacts=args.count)

    if leads:
        path = export_csv(leads)
        print(f"\nExported {len(leads)} leads to: {path}")

        if args.sync_attio:
            companies, people = sync_to_attio(leads)
            print(f"Synced to Attio: {companies} companies, {people} contacts")
    else:
        print("\nNo leads found. Try broadening your ICP criteria.")
