import sys
import json
import argparse
from perplexity import Perplexity

client = Perplexity()

SEC_DOMAINS = [
    "sec.gov",
    "edgar.sec.gov",
    "efts.sec.gov",
]


def search_sec_filings(query: str) -> dict:
    """Search SEC filings with domain-filtered web search."""
    response = client.responses.create(
        model="openai/gpt-5.4",
        input=query,
        tools=[{
            "type": "web_search",
            "filters": {
                "search_domain_filter": SEC_DOMAINS,
            },
        }],
        instructions=(
            "You are a financial analyst assistant. Search SEC filings for the requested "
            "information. Cite specific filing types (10-K, 10-Q, 8-K) and dates. "
            "Report exact numbers from the filings, not estimates."
        ),
        max_output_tokens=2048,
    )

    return {
        "query": query,
        "answer": response.output_text,
        "model": response.model,
        "cost": response.usage.cost.total_cost,
    }


def extract_financial_metrics(company: str, filing_type: str = "10-K") -> dict:
    """Extract structured financial metrics from SEC filings."""
    response = client.responses.create(
        model="openai/gpt-5.4",
        input=f"Find the most recent {filing_type} filing for {company} on SEC EDGAR and extract key financial metrics.",
        tools=[{
            "type": "web_search",
            "filters": {
                "search_domain_filter": SEC_DOMAINS,
            },
        }],
        instructions=(
            f"Search SEC EDGAR for {company}'s most recent {filing_type} filing. "
            "Extract the exact financial figures reported. Use numbers directly from the filing."
        ),
        response_format={
            "type": "json_schema",
            "json_schema": {
                "name": "sec_financials",
                "schema": {
                    "type": "object",
                    "properties": {
                        "company": {"type": "string"},
                        "ticker": {"type": "string"},
                        "filing_type": {"type": "string"},
                        "filing_period": {"type": "string"},
                        "filing_date": {"type": "string"},
                        "total_revenue": {"type": "string"},
                        "net_income": {"type": "string"},
                        "total_assets": {"type": "string"},
                        "total_debt": {"type": "string"},
                        "operating_income": {"type": "string"},
                        "eps_diluted": {"type": "string"},
                        "cash_and_equivalents": {"type": "string"},
                    },
                    "required": [
                        "company", "ticker", "filing_type", "filing_period", "filing_date",
                        "total_revenue", "net_income", "total_assets", "total_debt",
                        "operating_income", "eps_diluted", "cash_and_equivalents",
                    ],
                    "additionalProperties": False,
                },
            },
        },
    )

    return json.loads(response.output_text)


def compare_companies(tickers: list[str], metric: str = "revenue") -> list[dict]:
    """Compare a specific financial metric across multiple companies."""
    results = []
    for ticker in tickers:
        print(f"  Searching {ticker}...")
        try:
            data = extract_financial_metrics(ticker)
            results.append(data)
        except Exception as e:
            results.append({"ticker": ticker, "error": str(e)})
    return results


def search_filing_changes(company: str) -> str:
    """Search for material changes or risk factors in recent filings."""
    response = client.responses.create(
        model="openai/gpt-5.4",
        input=(
            f"What are the key risk factors and material changes disclosed in {company}'s "
            f"most recent 10-K or 10-Q filing on SEC EDGAR?"
        ),
        tools=[{
            "type": "web_search",
            "filters": {
                "search_domain_filter": SEC_DOMAINS,
            },
        }],
        instructions=(
            "Focus on risk factors (Item 1A) and material changes. "
            "Cite specific sections and filing dates."
        ),
        max_output_tokens=2048,
    )
    return response.output_text


def main():
    parser = argparse.ArgumentParser(description="SEC Filing Search")
    parser.add_argument("query", nargs="?", help="Search query for SEC filings")
    parser.add_argument("--extract", help="Extract financial metrics for a company ticker")
    parser.add_argument("--compare", nargs="+", help="Compare companies by ticker")
    parser.add_argument("--metric", default="revenue", help="Metric to compare")
    parser.add_argument("--risks", help="Search risk factors for a company")
    parser.add_argument("--json", action="store_true", help="Output as JSON")
    args = parser.parse_args()

    if args.compare:
        print(f"Comparing {len(args.compare)} companies...\n")
        results = compare_companies(args.compare, args.metric)
        if args.json:
            print(json.dumps(results, indent=2))
        else:
            for r in results:
                if "error" in r:
                    print(f"  {r['ticker']}: ERROR - {r['error']}")
                else:
                    print(f"  {r['company']} ({r['ticker']}) — {r['filing_type']} ({r['filing_period']})")
                    print(f"    Revenue: {r['total_revenue']}")
                    print(f"    Net Income: {r['net_income']}")
                    print(f"    Operating Income: {r['operating_income']}")
                    print(f"    EPS: {r['eps_diluted']}")
                    print()

    elif args.extract:
        print(f"Extracting financials for {args.extract}...\n")
        data = extract_financial_metrics(args.extract)
        if args.json:
            print(json.dumps(data, indent=2))
        else:
            print(f"{data['company']} ({data['ticker']})")
            print(f"Filing: {data['filing_type']} for {data['filing_period']} (filed {data['filing_date']})")
            print(f"  Revenue:          {data['total_revenue']}")
            print(f"  Net Income:       {data['net_income']}")
            print(f"  Operating Income: {data['operating_income']}")
            print(f"  Total Assets:     {data['total_assets']}")
            print(f"  Total Debt:       {data['total_debt']}")
            print(f"  Cash:             {data['cash_and_equivalents']}")
            print(f"  EPS (diluted):    {data['eps_diluted']}")

    elif args.risks:
        print(f"Searching risk factors for {args.risks}...\n")
        print(search_filing_changes(args.risks))

    elif args.query:
        result = search_sec_filings(args.query)
        if args.json:
            print(json.dumps(result, indent=2))
        else:
            print(result["answer"])

    else:
        print("Error: Provide a query, --extract TICKER, --compare TICKERS, or --risks COMPANY", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()