import sys
import requests

class SECClient:
    def __init__(self, user_agent):
        self.session = requests.Session()
        self.session.headers.update({
            "User-Agent": user_agent,
            "Accept-Encoding": "gzip, deflate",
        })

    def search_company(self, company_name):
        url = "https://www.sec.gov/files/company_tickers.json"
        try:
            response = self.session.get(url)
            response.raise_for_status()
            companies = response.json()
            company_name_lower = company_name.lower()
            for company_data in companies.values():
                if company_name_lower in company_data.get("title", "").lower():
                    return {
                        "cik": str(company_data.get("cik_str", "")).zfill(10),
                        "name": company_data.get("title", ""),
                        "ticker": company_data.get("ticker", "")
                    }
            return None
        except requests.RequestException as e:
            print(f"Error searching for company: {e}", file=sys.stderr)
            return None

    def get_company_filings(self, cik, form_types, from_date=None, to_date=None):
        url = f"https://data.sec.gov/submissions/CIK{cik}.json"
        try:
            response = self.session.get(url)
            response.raise_for_status()
            data = response.json()
            filings = data.get("filings", {}).get("recent", {})
            matches = []
            all_forms = filings.get("form", [])
            filing_dates = filings.get("filingDate", [])
            accession_numbers = filings.get("accessionNumber", [])
            primary_docs = filings.get("primaryDocument", [])
            for i, form_type in enumerate(all_forms):
                if form_type in form_types:
                    date = filing_dates[i]
                    if from_date and date < from_date:
                        continue
                    if to_date and date > to_date:
                        continue
                    matches.append({
                        "form": form_type,
                        "filingDate": date,
                        "accessionNumber": accession_numbers[i],
                        "accessionNumberClean": accession_numbers[i].replace("-", ""),
                        "primaryDocument": primary_docs[i],
                    })
            return matches
        except requests.RequestException as e:
            print(f"Error retrieving filings: {e}", file=sys.stderr)
            return []

    def get_filing_submission_text(self, cik, accession_number_clean, accession_number):
        try:
            url = f"https://www.sec.gov/Archives/edgar/data/{cik}/{accession_number_clean}/{accession_number}.txt"
            response = self.session.get(url)
            response.raise_for_status()
            return response.text
        except requests.RequestException as e:
            print(f"Error retrieving filing submission text: {e}", file=sys.stderr)
            return None

    def get_filing_document(self, cik, accession_number, document_name):
        try:
            url = f"https://www.sec.gov/Archives/edgar/data/{cik}/{accession_number}/{document_name}"
            response = self.session.get(url)
            response.raise_for_status()
            return response.text
        except requests.RequestException as e:
            print(f"Error downloading filing document: {e}", file=sys.stderr)
            return None
