import json
from pathlib import Path
from typing import Optional, Tuple

from src.negotiation_agent.Q_Table.domain.model.q_table import QTable
from src.negotiation_agent.Q_Table.domain.model.visit_table import VisitTable


class ModelRepository:
    """Simple repository to save/load Q-Table and VisitTable models."""

    def __init__(self, model_dir: Optional[str] = None):
        if model_dir is None:
            # Current: src/negotiation_agent/Q_Table/infra/repository/model_repository.py
            # Target Data: src/negotiation_agent/Q_Table/data/model
            # Path: ../../../data/model
            current_dir = Path(__file__).parent
            model_dir = current_dir.parent.parent / "data" / "model"
        
        self.model_dir = Path(model_dir)
        self.model_dir.mkdir(parents=True, exist_ok=True)
        self.q_table_path = self.model_dir / "q_table.json"
        self.visit_table_path = self.model_dir / "visit_table.json"

    def save(self, q_table: QTable, visit_table: VisitTable):
        with open(self.q_table_path, "w", encoding="utf-8") as f:
            json.dump(q_table.to_dict(), f, indent=2)
        
        with open(self.visit_table_path, "w", encoding="utf-8") as f:
            json.dump(visit_table.to_dict(), f, indent=2)
        
        print(f"Models saved to {self.model_dir}")

    def load(self) -> Tuple[Optional[QTable], Optional[VisitTable]]:
        q_table = None
        visit_table = None

        if self.q_table_path.exists():
            with open(self.q_table_path, "r", encoding="utf-8") as f:
                q_table = QTable.from_dict(json.load(f))
        
        if self.visit_table_path.exists():
            with open(self.visit_table_path, "r", encoding="utf-8") as f:
                visit_table = VisitTable.from_dict(json.load(f))
        
        return q_table, visit_table
