import sys
import json
import openai

REQUIRED_KEYS = {"company_name", "description"}

EXTRACTION_PROMPT = """\
Extract companies/entities acquired or owned by {target_company} from these SEC {form_type} excerpts.

For each, return the company name and a description including, if available:
- acquisition date
- full or partial ownership
- strategic rationale
- any rebranding, renaming, or integration into another product
- any subsequent divestitures

Return JSON: [{{"company_name": "...", "description": "..."}}]
Return [] if none found.

Excerpts:
{excerpts}"""

class LLMClient:
    def __init__(self, config):
        if not config.llm_api_key:
            raise ValueError("llm_api_key not set in ~/.acquiscan.yaml")
        self.client = openai.OpenAI(api_key=config.llm_api_key)
        self.model = config.llm_model

    def extract_acquisitions(self, excerpts, target_company, form_type):
        if not excerpts:
            return []

        prompt = EXTRACTION_PROMPT.format(
            target_company=target_company,
            form_type=form_type,
            excerpts="\n\n---\n\n".join(excerpts)
        )

        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
            )

            content = response.choices[0].message.content
            if not content:
                return []
            content = content.strip()
            if content.startswith("```"):
                lines = content.split("\n")
                content = "\n".join(lines[1:-1]) if len(lines) > 2 else content

            result = json.loads(content)
            if isinstance(result, dict) and "acquisitions" in result:
                result = result["acquisitions"]
            if not isinstance(result, list):
                return []

            return [item for item in result if isinstance(item, dict) and REQUIRED_KEYS <= item.keys()]
        except json.JSONDecodeError as e:
            print(f"  Error parsing LLM response: {e}", file=sys.stderr)
            return []
        except Exception as e:
            print(f"  Error calling LLM API: {e}", file=sys.stderr)
            return []
