"""Quick pipeline: Generate cascade chains directly from event metadata using LLM.

Skips GDELT/scraping (slow, rate-limited) and uses LLM + expert knowledge
to generate cascade chains from event metadata alone. GDELT-based pipeline
(02_collect_news.py + 03_extract_cascades.py) can be used later to augment
with real news data.
"""

import json
import logging
from datetime import date
from pathlib import Path

from src.data.cascade_extractor import save_cascade_chain, _extract_json, _infer_severity
from src.llm import create_llm_client
from src.llm.client import load_config
from src.models.schemas import CascadeChain, CascadeNode, FloodEvent

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


def generate_cascade_from_metadata(
    event: FloodEvent, llm_client, config: dict
) -> CascadeChain | None:
    """Generate a cascade chain using LLM + expert knowledge from event metadata."""
    variables = {
        "event_id": event.event_id,
        "country": event.country,
        "iso": event.iso,
        "location": event.location or event.country,
        "start_date": str(event.start_date),
        "origin": event.origin or "Unknown",
        "total_deaths": event.total_deaths or "Unknown",
        "total_affected": event.total_affected or "Unknown",
        "total_damage_k_usd": event.total_damage_k_usd or "Unknown",
    }

    response = llm_client.call_with_config(
        prompt_key="generate_cascade",
        knowledge_key="extraction",
        variables=variables,
        config=config,
    )

    try:
        json_str = _extract_json(response)
        data = json.loads(json_str)

        nodes = []
        for nd in data.get("cascade_events", []):
            nodes.append(CascadeNode(
                id=nd.get("id", f"E{len(nodes)+1}"),
                description=nd.get("description", ""),
                domain=nd.get("domain", "unknown"),
                severity=nd.get("severity", "medium"),
                time_offset_hours=nd.get("time_offset_hours"),
                mechanism=nd.get("mechanism", ""),
                parent_ids=nd.get("parent_ids", []),
            ))

        return CascadeChain(
            event_id=event.event_id,
            trigger_summary=data.get(
                "trigger_summary",
                f"Flood in {event.location or event.country} on {event.start_date}"
            ),
            trigger_country=event.country,
            trigger_iso=event.iso,
            trigger_date=event.start_date,
            trigger_severity=_infer_severity(event),
            cascade_events=nodes,
            source_articles=[],
            extraction_date=date.today(),
        )
    except (json.JSONDecodeError, KeyError, ValueError) as e:
        logger.error(f"Failed to parse LLM response for {event.event_id}: {e}")
        logger.error(f"Response: {response[:500]}")
        return None


def main():
    config = load_config()
    paths = config["paths"]

    # Load ALL events (both train and test — we want cascade chains for all)
    all_events_data = json.loads(Path(paths["events_catalog"]).read_text())
    all_events = [FloodEvent(**e) for e in all_events_data]
    logger.info(f"Loaded {len(all_events)} events total")

    llm_client = create_llm_client(config)
    chains_index = []

    for i, event in enumerate(all_events):
        logger.info(
            f"\n[{i+1}/{len(all_events)}] Generating cascade chain for: "
            f"{event.event_id} ({event.country}, {event.start_date})"
        )

        chain = generate_cascade_from_metadata(event, llm_client, config)
        if chain:
            path = save_cascade_chain(chain, config)
            logger.info(f"  Saved: {path} ({len(chain.cascade_events)} cascade nodes)")
            chains_index.append({
                "event_id": chain.event_id,
                "trigger_summary": chain.trigger_summary,
                "country": chain.trigger_country,
                "date": str(chain.trigger_date),
                "num_cascade_nodes": len(chain.cascade_events),
                "domains": list({n.domain for n in chain.cascade_events}),
                "file": str(path),
            })
        else:
            logger.warning(f"  Failed to generate chain for {event.event_id}")

    # Save index
    index_path = Path(paths["cascade_index"])
    index_path.parent.mkdir(parents=True, exist_ok=True)
    index_path.write_text(json.dumps(chains_index, indent=2, ensure_ascii=False))
    logger.info(f"\nSaved cascade chains index: {index_path}")
    logger.info(f"Total chains generated: {len(chains_index)}/{len(all_events)}")
    logger.info("Done. Run scripts/04_build_vectordb.py next.")


if __name__ == "__main__":
    main()